From 3a8073eda10cab0d302ce5946e24da2a323590d5 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Tue, 21 Jul 2026 21:10:54 -0700 Subject: [PATCH 1/3] Update [ghstack-poisoned] --- .../webgpu/runtime/WebGPUExecutionOptions.cpp | 10 +- .../webgpu/runtime/WebGPUExecutionOptions.h | 3 +- backends/webgpu/runtime/WebGPUGraph.cpp | 45 +- backends/webgpu/runtime/WebGPUGraph.h | 30 ++ backends/webgpu/runtime/WebGPUUtils.h | 98 ++++ .../ops/quantized_linear/QuantizedLinear.cpp | 278 ++++++---- backends/webgpu/runtime/ops/sdpa/Sdpa.cpp | 493 ++++++++++-------- .../ops/sdpa_fd_decode/SdpaFdDecode.cpp | 166 ++++-- .../runtime/ops/sdpa_fd_decode/SdpaFdDecode.h | 46 +- .../webgpu/test/native/test_dispatch_2d.cpp | 158 ++++++ .../webgpu/test/native/test_dynamic_shape.cpp | 299 +++++++++-- .../test/native/test_execution_options.cpp | 23 + .../test_dynamic_shape_export.py | 283 +++++++--- 13 files changed, 1468 insertions(+), 464 deletions(-) diff --git a/backends/webgpu/runtime/WebGPUExecutionOptions.cpp b/backends/webgpu/runtime/WebGPUExecutionOptions.cpp index 7362df98cfb..d3277fd1274 100644 --- a/backends/webgpu/runtime/WebGPUExecutionOptions.cpp +++ b/backends/webgpu/runtime/WebGPUExecutionOptions.cpp @@ -38,7 +38,12 @@ WebGPUExecutionPlan plan_webgpu_execution( size_t output_count, ExecuteConfig config, const std::vector& suppressible_outputs, - WebGPUGraphExecutionOptions options) { + WebGPUGraphExecutionOptions options, + const std::vector& enabled_dispatches) { + if (!enabled_dispatches.empty() && + enabled_dispatches.size() != dispatch_count) { + throw std::runtime_error("WebGPU: enabled dispatch count mismatch"); + } std::vector suppressed_dispatches(dispatch_count, false); std::vector copy_outputs(output_count, true); std::vector suppressed_outputs(output_count, false); @@ -73,7 +78,8 @@ WebGPUExecutionPlan plan_webgpu_execution( std::vector indices; indices.reserve(end - begin); for (size_t i = begin; i < end; i++) { - if (!suppressed_dispatches[i]) { + if (!suppressed_dispatches[i] && + (enabled_dispatches.empty() || enabled_dispatches[i])) { indices.push_back(i); } } diff --git a/backends/webgpu/runtime/WebGPUExecutionOptions.h b/backends/webgpu/runtime/WebGPUExecutionOptions.h index 691cd328e90..dab4723b7cc 100644 --- a/backends/webgpu/runtime/WebGPUExecutionOptions.h +++ b/backends/webgpu/runtime/WebGPUExecutionOptions.h @@ -49,7 +49,8 @@ WebGPUExecutionPlan plan_webgpu_execution( size_t output_count, ExecuteConfig config, const std::vector& suppressible_outputs, - WebGPUGraphExecutionOptions options); + WebGPUGraphExecutionOptions options, + const std::vector& enabled_dispatches = {}); WebGPUGraphExecutionOptions resolve_webgpu_graph_execution_options( const std::vector& delegate_outputs, diff --git a/backends/webgpu/runtime/WebGPUGraph.cpp b/backends/webgpu/runtime/WebGPUGraph.cpp index 5a567a36500..0dbcee2dbcc 100644 --- a/backends/webgpu/runtime/WebGPUGraph.cpp +++ b/backends/webgpu/runtime/WebGPUGraph.cpp @@ -436,6 +436,12 @@ void WebGPUGraph::build( if (!a) { continue; } + if (oc->name()->str() == "sym_size.int" && a->size() >= 3 && values) { + const auto* out = values->Get(a->Get(2)); + if (out && out->value_type() == vkgraph::GraphTypes::SymInt) { + dynamic_tensor_ids_.insert(static_cast(a->Get(0))); + } + } // f16 KV: tag sdpa K/V cache values (args[3],[4]) for half-size alloc. // Inert unless kv_f16_ (runtime opt-in) is set. if (kv_f16_ && a->size() > 4 && @@ -1733,26 +1739,47 @@ bool should_timestamp_query() { void WebGPUGraph::execute(const WebGPUGraphExecutionOptions& options) { const size_t n = dispatches_.size(); const size_t chunk = execute_config_.chunk_size; + std::vector enabled_dispatches(n, true); + for (size_t i = 0; i < n; i++) { + if (dispatches_[i].kind != WebGPUDispatch::Kind::Compute) { + continue; + } + const bool zero_x = dispatches_[i].workgroup_count_x == 0; + const bool zero_y = dispatches_[i].workgroup_count_y == 0; + if (zero_x != zero_y) { + throw std::runtime_error("WebGPU: dispatch has a half-zero grid"); + } + enabled_dispatches[i] = !zero_x; + } const WebGPUExecutionPlan plan = plan_webgpu_execution( n, output_copies_.size(), execute_config_, suppressible_outputs_, - options); + options, + enabled_dispatches); if (chunk == 0 || n <= chunk) { #ifdef WGPU_BACKEND_ENABLE_PROFILING + size_t active_compute_count = 0; + for (size_t i : plan.dispatch_chunks.front()) { + if (dispatches_[i].kind == WebGPUDispatch::Kind::Compute) { + active_compute_count++; + } + } // Bench: timestamp-query pool, null unless env-gated + feature present. WebGPUQueryPool* qp = nullptr; - if (should_timestamp_query() && n > 0) { + if (should_timestamp_query() && active_compute_count > 0) { if (auto* ctx = get_default_webgpu_context()) { if (ctx->timestamp_supported) { - if (!ctx->querypool || ctx->querypool->capacity() < n) { + if (!ctx->querypool || + ctx->querypool->capacity() < active_compute_count) { ctx->querypool = std::make_unique(); - ctx->querypool->initialize(device_, static_cast(n)); + ctx->querypool->initialize( + device_, static_cast(active_compute_count)); } qp = ctx->querypool.get(); - qp->reset(static_cast(n)); + qp->reset(static_cast(active_compute_count)); } } } @@ -1763,6 +1790,9 @@ void WebGPUGraph::execute(const WebGPUGraphExecutionOptions& options) { wgpuDeviceCreateCommandEncoder(device_, &enc_desc); // One pass per dispatch: enforces storage RAW ordering across deps. +#ifdef WGPU_BACKEND_ENABLE_PROFILING + uint32_t query_index = 0; +#endif for (size_t i : plan.dispatch_chunks.front()) { const auto& dispatch = dispatches_[i]; if (dispatch.kind == WebGPUDispatch::Kind::Copy) { @@ -1780,7 +1810,7 @@ void WebGPUGraph::execute(const WebGPUGraphExecutionOptions& options) { // tw must outlive BeginComputePass (the descriptor points at it). WGPUPassTimestampWrites tw = {}; if (qp) { - tw = qp->writes_for(static_cast(i)); + tw = qp->writes_for(query_index); pass_desc.timestampWrites = &tw; } #endif // WGPU_BACKEND_ENABLE_PROFILING @@ -1796,10 +1826,11 @@ void WebGPUGraph::execute(const WebGPUGraphExecutionOptions& options) { #ifdef WGPU_BACKEND_ENABLE_PROFILING if (qp) { qp->record( - static_cast(i), + query_index, dispatch.kernel_name, {dispatch.workgroup_count_x, dispatch.workgroup_count_y, 1}, {1, 1, 1}); + query_index++; } #endif // WGPU_BACKEND_ENABLE_PROFILING } diff --git a/backends/webgpu/runtime/WebGPUGraph.h b/backends/webgpu/runtime/WebGPUGraph.h index c520cb06f84..c12d1e65e0f 100644 --- a/backends/webgpu/runtime/WebGPUGraph.h +++ b/backends/webgpu/runtime/WebGPUGraph.h @@ -18,6 +18,7 @@ #include #include +#include #include namespace executorch::backends::webgpu { @@ -189,6 +190,14 @@ class WebGPUGraph { symint_dim_sources_.push_back({symint_id, tensor_id, dim}); } + bool tensor_has_dynamic_dims(int tensor_id) const { + return dynamic_tensor_ids_.count(tensor_id) != 0; + } + + bool has_dynamic_shapes() const { + return !dynamic_tensor_ids_.empty(); + } + // Execute-time select_as_symint read; mirrors Vulkan select_as_symint_impl. void update_symints_from_inputs(const std::vector& inputs); @@ -225,6 +234,25 @@ class WebGPUGraph { return dispatches_.size(); } + size_t register_dispatch_route_group( + const std::vector& ranges) { + return dispatch_routes_.register_group( + dispatches_.size(), ranges, [&](size_t i) { + return dispatches_[i].kind == WebGPUDispatch::Kind::Compute; + }); + } + + void select_dispatch_route( + size_t group, + size_t active_route, + const std::vector& active_grids) { + dispatch_routes_.select( + group, active_route, active_grids, [&](size_t i, utils::WgCount grid) { + dispatches_[i].workgroup_count_x = grid.x; + dispatches_[i].workgroup_count_y = grid.y; + }); + } + WGPUDevice device() const { return device_; } @@ -383,6 +411,7 @@ class WebGPUGraph { std::unordered_map symints_; std::vector symint_sources_; std::vector symint_dim_sources_; + std::unordered_set dynamic_tensor_ids_; // Resize hooks + the set of SymInts changed since the last propagate_resize. struct ResizeHook { @@ -434,6 +463,7 @@ class WebGPUGraph { std::vector suppressible_outputs_; std::vector dispatches_; + utils::DispatchRouteRegistry dispatch_routes_; // Prepack-routed constant sources (offset/named-key + size); the prepack node // materializes these once. constant_data_/named_data_map_ point at the .pte diff --git a/backends/webgpu/runtime/WebGPUUtils.h b/backends/webgpu/runtime/WebGPUUtils.h index b41b02bf6e4..da37c7c0438 100644 --- a/backends/webgpu/runtime/WebGPUUtils.h +++ b/backends/webgpu/runtime/WebGPUUtils.h @@ -12,6 +12,7 @@ #include #include +#include #include #include #include @@ -63,6 +64,103 @@ struct WgCount { uint32_t y; }; +struct DispatchRange { + size_t begin; + size_t end; +}; + +constexpr bool should_record_q4gsw_dual_route( + uint32_t max_m, + bool bicol_eligible, + bool has_dynamic_shapes) { + return max_m > 1u && bicol_eligible && has_dynamic_shapes; +} + +constexpr bool should_record_sdpa_dual_route( + bool fd_eligible, + bool has_dynamic_sequence) { + return fd_eligible && has_dynamic_sequence; +} + +class DispatchRouteRegistry { + public: + template + size_t register_group( + size_t dispatch_count, + const std::vector& ranges, + IsCompute&& is_compute) { + if (dispatch_count < owners_.size() || ranges.size() < 2) { + throw std::runtime_error("invalid WebGPU dispatch route group"); + } + + std::vector claimed(dispatch_count, false); + for (const auto& range : ranges) { + if (range.begin >= range.end || range.end > dispatch_count) { + throw std::runtime_error("invalid WebGPU dispatch route range"); + } + for (size_t i = range.begin; i < range.end; i++) { + if (!is_compute(i)) { + throw std::runtime_error( + "WebGPU dispatch route contains a copy command"); + } + if (claimed[i] || (i < owners_.size() && owners_[i] != kNoOwner)) { + throw std::runtime_error("overlapping WebGPU dispatch route ranges"); + } + claimed[i] = true; + } + } + + const size_t group = groups_.size(); + owners_.resize(dispatch_count, kNoOwner); + for (size_t i = 0; i < claimed.size(); i++) { + if (claimed[i]) { + owners_[i] = group; + } + } + groups_.push_back(ranges); + return group; + } + + template + void select( + size_t group, + size_t active_route, + const std::vector& active_grids, + SetGrid&& set_grid) const { + if (group >= groups_.size()) { + throw std::runtime_error("invalid WebGPU dispatch route group"); + } + const auto& ranges = groups_[group]; + if (active_route >= ranges.size()) { + throw std::runtime_error("invalid active WebGPU dispatch route"); + } + const auto& active = ranges[active_route]; + if (active_grids.size() != active.end - active.begin) { + throw std::runtime_error("WebGPU dispatch route grid count mismatch"); + } + for (const auto& grid : active_grids) { + if (grid.x == 0 || grid.y == 0) { + throw std::runtime_error( + "active WebGPU dispatch route has a zero grid"); + } + } + + for (const auto& range : ranges) { + for (size_t i = range.begin; i < range.end; i++) { + set_grid(i, {0, 0}); + } + } + for (size_t i = 0; i < active_grids.size(); i++) { + set_grid(active.begin + i, active_grids[i]); + } + } + + private: + static constexpr size_t kNoOwner = static_cast(-1); + std::vector> groups_; + std::vector owners_; +}; + // Device's max workgroups per dispatch dimension; the WebGPU spec-default floor // (65535) if the query fails — never under-reports a real device's capacity. inline uint32_t queried_max_workgroups(WGPUDevice device) { diff --git a/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp b/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp index a931e22b08d..ffeeffae5d1 100644 --- a/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp +++ b/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp @@ -162,6 +162,75 @@ uint32_t compute_q4gsw_workgroup_count( device, static_cast(total_tiles), wg_size, op_name); } +struct Q4gswExecutionState { + Q4gswParams params; + std::vector output_dims; + size_t active_route; + utils::WgCount active_grid; +}; + +constexpr size_t kQ4gswBicolRoute = 0; +constexpr size_t kQ4gswPrefillRoute = 1; + +Q4gswExecutionState make_q4gsw_execution_state( + WGPUDevice device, + const std::vector& input_dims, + uint32_t max_m, + uint32_t K, + uint32_t N, + uint32_t K_packed, + uint32_t gs, + uint32_t padded_N, + uint32_t has_bias, + uint32_t wg_size, + bool use_single_gemv, + bool use_dual_route, + bool prefill_use_steel, + bool prefill_use_shmem_gemm) { + if (input_dims.empty()) { + throw std::runtime_error("WebGPU linear_q4gsw(resize): empty input dims"); + } + const uint64_t numel = utils::numel_of(input_dims); + if (numel % static_cast(K) != 0u) { + throw std::runtime_error( + "WebGPU linear_q4gsw(resize): live input numel not a multiple of K"); + } + const uint64_t live_m = numel / static_cast(K); + if (live_m == 0u) { + throw std::runtime_error("WebGPU linear_q4gsw(resize): live M == 0"); + } + if (live_m > max_m) { + throw std::runtime_error( + "WebGPU linear_q4gsw(resize): live M exceeds the build-time max"); + } + const uint32_t m = static_cast(live_m); + const bool use_gemv = use_single_gemv || (use_dual_route && m == 1u); + const uint32_t workgroup_count = compute_q4gsw_workgroup_count( + device, + use_gemv, + !use_gemv && prefill_use_steel, + !use_gemv && prefill_use_shmem_gemm, + m, + N, + wg_size, + "linear_q4gsw(resize)"); + + Q4gswExecutionState state = {}; + state.params.M = m; + state.params.N = N; + state.params.K = K; + state.params.K_packed = K_packed; + state.params.group_size = gs; + state.params.padded_N = padded_N; + state.params.has_bias = has_bias; + state.output_dims = input_dims; + state.output_dims.back() = static_cast(N); + state.active_route = + use_dual_route ? (use_gemv ? kQ4gswBicolRoute : kQ4gswPrefillRoute) : 0u; + state.active_grid = {workgroup_count, 1u}; + return state; +} + // et_vk.linear_q4gsw args: [in, weight, scales, group_size, bias, out]. void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { const int in_id = args.at(0); @@ -243,9 +312,12 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { // M==1 -> bicol GEMV; M>1 -> steel GEMM (preferred) else shmem else tiled. const uint32_t wg_size = utils::clamp_workgroup_size(device, kQ4gswLinearWorkgroupSizeX); - const bool use_gemv = (M == 1u && K % 8u == 0u && gs % 8u == 0u); + const bool bicol_eligible = K % 8u == 0u && gs % 8u == 0u; + const bool use_gemv = M == 1u && bicol_eligible; + const bool use_dual_route = utils::should_record_q4gsw_dual_route( + M, bicol_eligible, graph.has_dynamic_shapes()); // GEMV (bicol) is a pow2 tree reduction; compute its size only when used. - const uint32_t gemv_wg_size = use_gemv + const uint32_t gemv_wg_size = (use_gemv || use_dual_route) ? utils::clamp_workgroup_size_pow2( device, kQ4gswLinearCoop4BicolWorkgroupSizeX) : 0u; @@ -258,6 +330,9 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { // large K/N thresholds; otherwise the register-tiled path handles it. const bool use_shmem_gemm = !use_gemv && !use_steel && (K >= kQ4gswShmemMinDim || N >= kQ4gswShmemNMinDim); + const char* prefill_shader_src = use_steel ? kQ4gswLinearGemmSteelWGSL + : use_shmem_gemm ? kQ4gswLinearGemmShmemWGSL + : kQ4gswLinearWGSL; const char* shader_src = use_gemv ? kQ4gswLinearCoop4BicolWGSL : use_steel ? kQ4gswLinearGemmSteelWGSL : use_shmem_gemm ? kQ4gswLinearGemmShmemWGSL @@ -271,9 +346,10 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { // each u32 weight word once + hoists the per-column scale (half re-reads // them ~8x/~16x). Needs group_size % BK == 0 so the hoisted scale is // constant across the BK tile; else the per-nibble `half` kernel. - shader_src = (gs % kQ4gswSteelBK == 0u) + prefill_shader_src = (gs % kQ4gswSteelBK == 0u) ? kQ4gswLinearGemmSteelHalfPwdqWGSL : kQ4gswLinearGemmSteelHalfWGSL; + shader_src = prefill_shader_src; } } // f16-accumulate: pwdq staging with an f16 register accumulator. @@ -284,18 +360,10 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { if (use_steel && graph.f16_accumulate_gemm() && (gs % kQ4gswSteelBK == 0u)) { const WebGPUContext* ctx = get_default_webgpu_context(); if (ctx != nullptr && ctx->shader_f16_supported) { - shader_src = kQ4gswLinearGemmSteelHalfPwdqF16accWGSL; + prefill_shader_src = kQ4gswLinearGemmSteelHalfPwdqF16accWGSL; + shader_src = prefill_shader_src; } } - const uint32_t workgroup_count = compute_q4gsw_workgroup_count( - device, - use_gemv, - use_steel, - use_shmem_gemm, - M, - N, - wg_size, - "linear_q4gsw"); // Optional bias: real buffer if present, else a dummy for the fixed layout. uint32_t has_bias = 0; @@ -315,14 +383,21 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { bias_buffer = graph.create_scratch_buffer(4); } - Q4gswParams params = {}; - params.M = M; - params.N = N; - params.K = K; - params.K_packed = K_packed; - params.group_size = gs; - params.padded_N = padded_N; - params.has_bias = has_bias; + const Q4gswExecutionState initial_state = make_q4gsw_execution_state( + device, + in.dims, + M, + K, + N, + K_packed, + gs, + padded_N, + has_bias, + wg_size, + use_gemv, + use_dual_route, + use_steel, + use_shmem_gemm); WGPUBufferDescriptor uniform_desc = {}; uniform_desc.size = sizeof(Q4gswParams); @@ -331,17 +406,10 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { WGPUBuffer uniform_buffer = wgpuDeviceCreateBuffer(device, &uniform_desc); void* mapped = wgpuBufferGetMappedRange(uniform_buffer, 0, sizeof(Q4gswParams)); - std::memcpy(mapped, ¶ms, sizeof(Q4gswParams)); + std::memcpy(mapped, &initial_state.params, sizeof(Q4gswParams)); wgpuBufferUnmap(uniform_buffer); graph.add_uniform_buffer_bytes(sizeof(Q4gswParams)); - WGPUShaderSourceWGSL wgsl_desc = {}; - wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; - wgsl_desc.code = {shader_src, WGPU_STRLEN}; - WGPUShaderModuleDescriptor shader_desc = {}; - shader_desc.nextInChain = &wgsl_desc.chain; - WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc); - // Bind group layout: out (rw) + in/weight/scales/bias (ro storage) + uniform. WGPUBindGroupLayoutEntry entries[6] = {}; entries[0].binding = 0; @@ -367,22 +435,6 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { WGPUPipelineLayout pipeline_layout = wgpuDeviceCreatePipelineLayout(device, &pl_desc); - // GEMV/tiled wire an override wg_size; steel (256) + shmem (64) are fixed. - const bool fixed_wg = use_steel || use_shmem_gemm; - WGPUConstantEntry wg_size_constant = {}; - wg_size_constant.key = {"wg_size", WGPU_STRLEN}; - wg_size_constant.value = - static_cast(use_gemv ? gemv_wg_size : wg_size); - - WGPUComputePipelineDescriptor pipeline_desc = {}; - pipeline_desc.layout = pipeline_layout; - pipeline_desc.compute.module = shader; - pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; - pipeline_desc.compute.constantCount = fixed_wg ? 0u : 1u; - pipeline_desc.compute.constants = fixed_wg ? nullptr : &wg_size_constant; - WGPUComputePipeline pipeline = - wgpuDeviceCreateComputePipeline(device, &pipeline_desc); - WGPUBindGroupEntry bg_entries[6] = {}; bg_entries[0].binding = 0; bg_entries[0].buffer = out.buffer; @@ -409,12 +461,72 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { bg_desc.entries = bg_entries; WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc); - const size_t dispatch_idx = graph.add_dispatch( - {pipeline, bind_group, workgroup_count, "linear_q4gsw"}); + auto make_pipeline = [&](const char* source, + bool fixed_wg, + uint32_t override_wg_size) { + WGPUShaderSourceWGSL wgsl_desc = {}; + wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL; + wgsl_desc.code = {source, WGPU_STRLEN}; + WGPUShaderModuleDescriptor shader_desc = {}; + shader_desc.nextInChain = &wgsl_desc.chain; + WGPUShaderModule shader = + wgpuDeviceCreateShaderModule(device, &shader_desc); + + WGPUConstantEntry wg_size_constant = {}; + wg_size_constant.key = {"wg_size", WGPU_STRLEN}; + wg_size_constant.value = static_cast(override_wg_size); + WGPUComputePipelineDescriptor pipeline_desc = {}; + pipeline_desc.layout = pipeline_layout; + pipeline_desc.compute.module = shader; + pipeline_desc.compute.entryPoint = {"main", WGPU_STRLEN}; + pipeline_desc.compute.constantCount = fixed_wg ? 0u : 1u; + pipeline_desc.compute.constants = fixed_wg ? nullptr : &wg_size_constant; + WGPUComputePipeline pipeline = + wgpuDeviceCreateComputePipeline(device, &pipeline_desc); + wgpuShaderModuleRelease(shader); + return pipeline; + }; + + const bool fixed_prefill_wg = use_steel || use_shmem_gemm; + const char* prefill_label = use_steel ? "linear_q4gsw_steel" + : use_shmem_gemm ? "linear_q4gsw_shmem" + : "linear_q4gsw_tiled"; + size_t dispatch_idx = 0; + size_t route_group = 0; + if (use_dual_route) { + // Each recorded dispatch owns one bind-group reference. + wgpuBindGroupAddRef(bind_group); + WGPUComputePipeline bicol_pipeline = + make_pipeline(kQ4gswLinearCoop4BicolWGSL, false, gemv_wg_size); + const size_t bicol_idx = graph.add_dispatch( + {bicol_pipeline, + bind_group, + initial_state.active_grid.x, + "linear_q4gsw_coop4_bicol"}); + WGPUComputePipeline prefill_pipeline = + make_pipeline(prefill_shader_src, fixed_prefill_wg, wg_size); + const size_t prefill_idx = graph.add_dispatch( + {prefill_pipeline, + bind_group, + initial_state.active_grid.x, + prefill_label}); + route_group = graph.register_dispatch_route_group( + {{bicol_idx, bicol_idx + 1}, {prefill_idx, prefill_idx + 1}}); + graph.select_dispatch_route( + route_group, initial_state.active_route, {initial_state.active_grid}); + } else { + const bool fixed_wg = use_gemv ? false : fixed_prefill_wg; + WGPUComputePipeline pipeline = + make_pipeline(shader_src, fixed_wg, use_gemv ? gemv_wg_size : wg_size); + dispatch_idx = graph.add_dispatch( + {pipeline, + bind_group, + initial_state.active_grid.x, + use_gemv ? "linear_q4gsw_coop4_bicol" : prefill_label}); + } - // Dynamic shapes: recompute dispatch + params.M for the live M. use_gemv and - // use_shmem_gemm are captured (routing is fixed at build); the helper re-runs - // the same path's workgroup-count formula with the live m. + // Dynamic shapes: recompute one shared Params block and select exactly one + // writer. The prefill pipeline remains the route chosen from max M. graph.add_tensor_resize_hook( in_id, [in_id, @@ -428,57 +540,41 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { has_bias, wg_size, use_gemv, + use_dual_route, use_steel, use_shmem_gemm, dispatch_idx, + route_group, uniform_buffer](WebGPUGraph& g) { const auto& d = g.cur_dims(in_id); - if (d.empty()) { - throw std::runtime_error( - "WebGPU linear_q4gsw(resize): empty input dims"); - } - const uint64_t numel = utils::numel_of(d); - if (numel % static_cast(K) != 0u) { - throw std::runtime_error( - "WebGPU linear_q4gsw(resize): live input numel not a multiple " - "of K"); - } - const uint32_t m = - static_cast(numel / static_cast(K)); - if (m == 0u) { - throw std::runtime_error("WebGPU linear_q4gsw(resize): live M == 0"); - } - // Buffers/bind-groups were sized for the build-time max M; a larger - // live M would write out of bounds. - if (m > M) { - throw std::runtime_error( - "WebGPU linear_q4gsw(resize): live M exceeds the build-time max"); - } - const uint32_t wgc = compute_q4gsw_workgroup_count( + const Q4gswExecutionState state = make_q4gsw_execution_state( g.device(), - use_gemv, - use_steel, - use_shmem_gemm, - m, + d, + M, + K, N, + K_packed, + gs, + padded_N, + has_bias, wg_size, - "linear_q4gsw(resize)"); - Q4gswParams p = {}; - p.M = m; - p.N = N; - p.K = K; - p.K_packed = K_packed; - p.group_size = gs; - p.padded_N = padded_N; - p.has_bias = has_bias; - wgpuQueueWriteBuffer(g.queue(), uniform_buffer, 0, &p, sizeof(p)); - g.dispatch_at(dispatch_idx).workgroup_count_x = wgc; - std::vector od(d.begin(), d.end()); - od.back() = static_cast(N); - g.set_cur_dims(out_id, od); + use_gemv, + use_dual_route, + use_steel, + use_shmem_gemm); + wgpuQueueWriteBuffer( + g.queue(), uniform_buffer, 0, &state.params, sizeof(state.params)); + if (use_dual_route) { + g.select_dispatch_route( + route_group, state.active_route, {state.active_grid}); + } else { + auto& dispatch = g.dispatch_at(dispatch_idx); + dispatch.workgroup_count_x = state.active_grid.x; + dispatch.workgroup_count_y = state.active_grid.y; + } + g.set_cur_dims(out_id, state.output_dims); }); - wgpuShaderModuleRelease(shader); wgpuBindGroupLayoutRelease(bgl); wgpuPipelineLayoutRelease(pipeline_layout); // Graph owns it so the resize hook can rewrite it; freed in the dtor. diff --git a/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp b/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp index f3345c89c77..817f96ec19d 100644 --- a/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp +++ b/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp @@ -74,6 +74,22 @@ struct ComputeOutParams { }; static_assert(sizeof(ComputeOutParams) == 32, "ComputeOutParams must be 32B"); +struct SdpaLiveState { + int64_t s; + int64_t pos; + int64_t context_len; + UpdateCacheParams update_cache; + AttnWeightsParams attn_weights; + SoftmaxParams softmax; + ComputeOutParams compute_out; + utils::WgCount update_cache_grid; + utils::WgCount qk_grid; + utils::WgCount softmax_grid; + utils::WgCount av_grid; + bool use_fd; + SdpaFdDecodeState fd; +}; + // Param-struct builder helpers — used in both initial build and resize hook. static UpdateCacheParams make_update_cache_params( uint64_t kv_numel, @@ -419,20 +435,6 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { throw std::runtime_error("WebGPU sdpa: only is_causal=true is supported"); } - // KV cache written in place; only attn_weights/softmax need scratch. - const uint64_t aw_floats = static_cast(Hq) * - static_cast(S) * static_cast(context_len); - // Dynamic input_pos: size+bind scratch for Cmax (no realloc; covers any ctx). - const uint64_t aw_cap_floats = static_cast(Hq) * - static_cast(S) * - static_cast(dynamic_pos ? Cmax : context_len); - const uint64_t aw_bytes = aw_cap_floats * sizeof(float); - - // Dynamic input_pos: the resize hook rewrites these per step. - WGPUBuffer uc_k_buf = nullptr, uc_v_buf = nullptr, qk_buf = nullptr, - softmax_buf = nullptr, av_buf = nullptr; - size_t qk_idx = 0, uc_k_idx = 0, uc_v_idx = 0, softmax_idx = 0, av_idx = 0; - const WGPUDevice device = graph.device(); const uint32_t uc_wg = utils::clamp_workgroup_size(device, kUpdateCacheWorkgroupSizeX); @@ -440,258 +442,329 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { 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); + const bool fd_eligible = D <= kSdpaFdMaxHeadDim; + const int64_t pos_const = input_pos; + + auto compute_live_state = [q_id, + k_id, + v_id, + out_id, + qn, + kn, + S, + dynamic_pos, + input_pos_id, + pos_const, + Hq, + Hkv, + D, + Cmax, + g, + scale, + uc_wg, + qk_wg, + av_wg, + fd_eligible](WebGPUGraph& gr) { + SdpaLiveState state = {}; + const auto& q_live_dims = gr.cur_dims(q_id); + state.s = q_live_dims[qn - 3]; + state.pos = dynamic_pos ? static_cast(gr.read_symint(input_pos_id)) + : pos_const; + if (state.s <= 0 || state.pos < 0 || state.s > S) { + throw std::runtime_error("WebGPU sdpa: invalid live S or input_pos"); + } + if (gr.cur_dims(k_id)[kn - 3] != state.s || + gr.cur_dims(v_id)[gr.cur_dims(v_id).size() - 3] != state.s) { + throw std::runtime_error("WebGPU sdpa: live q/k/v seq_len mismatch"); + } + const auto& out_max_dims = gr.get_tensor(out_id).dims; + if (out_max_dims.size() != q_live_dims.size()) { + throw std::runtime_error("WebGPU sdpa: output rank must match q"); + } + for (size_t i = 0; i < q_live_dims.size(); i++) { + if (q_live_dims[i] <= 0 || q_live_dims[i] > out_max_dims[i]) { + throw std::runtime_error( + "WebGPU sdpa: live output shape exceeds allocation"); + } + } + state.context_len = state.s + state.pos; + if (state.context_len <= 0 || state.context_len > Cmax || + state.s > UINT32_MAX || state.pos > UINT32_MAX || + state.context_len > UINT32_MAX) { + throw std::runtime_error( + "WebGPU sdpa: live dimensions exceed cache or uint32 capacity"); + } + + const uint64_t kv_numel = static_cast(state.s) * + static_cast(Hkv) * static_cast(D); + const uint64_t kv_offset = static_cast(state.pos) * + 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) { + throw std::runtime_error("WebGPU sdpa: live workload exceeds uint32"); + } + + state.update_cache = make_update_cache_params( + kv_numel, static_cast(kv_offset), cache_numel); + state.attn_weights = make_attn_weights_params( + state.s, Hq, Hkv, D, state.context_len, state.pos, g, scale); + 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); + 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)"); + state.use_fd = fd_eligible && state.s == 1; + if (fd_eligible) { + state.fd = make_sdpa_fd_decode_state( + gr.device(), Hq, Hkv, D, state.context_len, g, scale); + } + return state; + }; + + 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_bytes = aw_cap_floats * sizeof(float); - // Dispatches 1-2: write new K/V into the caches (reuses update_cache). - const uint64_t kv_numel = static_cast(S) * - static_cast(Hkv) * static_cast(D); - const uint32_t kv_dst_offset = static_cast( - static_cast(input_pos) * static_cast(Hkv) * - static_cast(D)); - uc_k_buf = record_update_cache_dispatch( + WGPUBuffer uc_k_buf = record_update_cache_dispatch( graph, device, k_cache, k, - kv_numel, - kv_dst_offset, - numel(k_cache), + initial_state.update_cache.numel, + initial_state.update_cache.dst_offset, + initial_state.update_cache.cache_numel, uc_wg, true, "update_cache(K)"); - uc_v_buf = record_update_cache_dispatch( + WGPUBuffer uc_v_buf = record_update_cache_dispatch( graph, device, v_cache, v, - kv_numel, - kv_dst_offset, - numel(v_cache), + initial_state.update_cache.numel, + initial_state.update_cache.dst_offset, + initial_state.update_cache.cache_numel, uc_wg, true, "update_cache(V)"); - uc_k_idx = graph.num_dispatches() - 2; - uc_v_idx = graph.num_dispatches() - 1; - - // FlashDecoding decode (S==1, static pos). Shapes FD can't handle (head dim - // > kSdpaFdMaxHeadDim) fall through to the materialized path below. - if (S == 1 && !dynamic_pos && D <= kSdpaFdMaxHeadDim) { - sdpa_fd_decode_dispatch( - graph, q, k_cache, v_cache, out, Hq, Hkv, D, context_len, g, scale); - return; - } - - // QK/softmax scratch — allocated only on the non-FD path (Hq*S*Cmax prefill). - WGPUBuffer attn_weights = graph.acquire_scratch(aw_bytes); - WebGPUGraph::ScopedScratch attn_weights_guard(&graph, attn_weights); - WGPUBuffer attn_weights_softmax = graph.acquire_scratch(aw_bytes); - WebGPUGraph::ScopedScratch attn_weights_softmax_guard( - &graph, attn_weights_softmax); - - // --- Dispatch 3: QK -> attn_weights. One thread per TM x TN tile. - { - if (aw_floats > UINT32_MAX) { - throw std::runtime_error( - "WebGPU sdpa: Hq*S*context_len exceeds uint32 max"); - } - const int64_t qk_tiles = Hq * utils::div_up(S, kSdpaTileM) * - utils::div_up(context_len, kSdpaTileN); - const utils::WgCount wgc = utils::compute_2d_workgroup_count( - device, static_cast(qk_tiles), qk_wg, "QK"); - AttnWeightsParams p = make_attn_weights_params( - S, Hq, Hkv, D, context_len, input_pos, g, scale); - WGPUBuffer ubuf = graph.make_uniform_buffer(&p, sizeof(p)); - BufferBinding bindings[3] = { + const size_t uc_k_idx = graph.num_dispatches() - 2; + const size_t uc_v_idx = graph.num_dispatches() - 1; + 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 record_materialized = dual_route || !initial_state.use_fd; + const bool record_fd = dual_route || initial_state.use_fd; + + WGPUBuffer qk_buf = nullptr; + WGPUBuffer softmax_buf = nullptr; + WGPUBuffer av_buf = nullptr; + size_t qk_idx = 0; + size_t softmax_idx = 0; + size_t av_idx = 0; + utils::DispatchRange materialized_range = {}; + if (record_materialized) { + WGPUBuffer attn_weights = graph.acquire_scratch(aw_bytes); + WebGPUGraph::ScopedScratch attn_weights_guard(&graph, attn_weights); + WGPUBuffer attn_weights_softmax = graph.acquire_scratch(aw_bytes); + WebGPUGraph::ScopedScratch attn_weights_softmax_guard( + &graph, attn_weights_softmax); + + 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 = kSdpaComputeAttnWeightsWGSL; - if (graph.kv_f16()) { - qk_src = kSdpaComputeAttnWeightsHalfWGSL; - } + const char* qk_src = graph.kv_f16() ? kSdpaComputeAttnWeightsHalfWGSL + : kSdpaComputeAttnWeightsWGSL; build_dispatch( graph, qk_src, - bindings, + qk_bindings, 3, - ubuf, - sizeof(p), - wgc.x, - wgc.y, + qk_buf, + sizeof(AttnWeightsParams), + initial_state.qk_grid.x, + initial_state.qk_grid.y, qk_wg, true, "sdpa_compute_attn_weights"); - qk_buf = ubuf; qk_idx = graph.num_dispatches() - 1; - } - // Dispatch 4: softmax, one workgroup per (h,s) row of width context_len. - { - // One workgroup per (h,s) row; wg_size 1 keeps the device dispatch check. - const utils::WgCount wgc = utils::compute_2d_workgroup_count( - device, static_cast(Hq * S), 1, "softmax"); - const uint32_t sm_wg = - utils::clamp_workgroup_size_pow2(device, kSdpaSoftmaxWorkgroupSizeX); - SoftmaxParams p = make_softmax_params(Hq, S, context_len); - WGPUBuffer ubuf = graph.make_uniform_buffer(&p, sizeof(p)); - BufferBinding bindings[2] = { + 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( graph, kSdpaSoftmaxWGSL, - bindings, + softmax_bindings, 2, - ubuf, - sizeof(p), - wgc.x, - wgc.y, + softmax_buf, + sizeof(SoftmaxParams), + initial_state.softmax_grid.x, + initial_state.softmax_grid.y, sm_wg, true, "sdpa_softmax"); - softmax_buf = ubuf; softmax_idx = graph.num_dispatches() - 1; - } - // --- Dispatch 5: AV -> out. One thread per TM x TN tile. - { - const int64_t av_tiles = - Hq * utils::div_up(S, kSdpaTileM) * utils::div_up(D, kSdpaTileN); - const utils::WgCount wgc = utils::compute_2d_workgroup_count( - device, static_cast(av_tiles), av_wg, "AV"); - ComputeOutParams p = make_compute_out_params(S, Hq, Hkv, D, context_len, g); - WGPUBuffer ubuf = graph.make_uniform_buffer(&p, sizeof(p)); - BufferBinding bindings[3] = { + 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 = kSdpaComputeOutWGSL; - if (graph.kv_f16()) { - av_src = kSdpaComputeOutHalfWGSL; - } + const char* av_src = + graph.kv_f16() ? kSdpaComputeOutHalfWGSL : kSdpaComputeOutWGSL; build_dispatch( graph, av_src, - bindings, + av_bindings, 3, - ubuf, - sizeof(p), - wgc.x, - wgc.y, + av_buf, + sizeof(ComputeOutParams), + initial_state.av_grid.x, + initial_state.av_grid.y, av_wg, true, "sdpa_compute_out"); - av_buf = ubuf; av_idx = graph.num_dispatches() - 1; + materialized_range.end = graph.num_dispatches(); } - // Per-step recompute: live S (q resize) or input_pos (SymInt); inert if - // static. - const int64_t pos_const = input_pos; - auto sdpa_resize = [q_id, - qn, - S, - out_id, - dynamic_pos, - input_pos_id, - pos_const, - Hq, - Hkv, - D, - Cmax, - g, - scale, - qk_idx, - uc_k_idx, - uc_v_idx, - softmax_idx, - av_idx, - uc_wg, - qk_wg, - av_wg, - uc_k_buf, - uc_v_buf, - qk_buf, - softmax_buf, - av_buf](WebGPUGraph& gr) { - const int64_t s = gr.cur_dims(q_id)[qn - 3]; - const int64_t pos = dynamic_pos - ? static_cast(gr.read_symint(input_pos_id)) - : pos_const; - if (s <= 0 || pos < 0) { - throw std::runtime_error("WebGPU sdpa: invalid live S or input_pos"); - } - // Scratch (attn_weights/softmax) is sized at build for S=max; a larger live - // S would overrun it. Make that invariant load-bearing. - if (s > S) { - throw std::runtime_error( - "WebGPU sdpa: live S exceeds the build-time max (scratch capacity)"); - } - const int64_t ctx = s + pos; - if (ctx <= 0 || ctx > Cmax) { - throw std::runtime_error( - "WebGPU sdpa: context_len exceeds cache capacity"); + SdpaFdDecodeResources fd_resources = {}; + size_t route_group = 0; + if (record_fd) { + fd_resources = record_sdpa_fd_decode_dispatches( + graph, q, k_cache, v_cache, out, initial_state.fd); + } + if (dual_route) { + route_group = graph.register_dispatch_route_group( + {materialized_range, fd_resources.dispatch_range}); + } + + auto refresh_state = [compute_live_state, + q_id, + out_id, + dual_route, + record_materialized, + record_fd, + fixed_use_fd = initial_state.use_fd, + route_group, + uc_k_idx, + uc_v_idx, + qk_idx, + softmax_idx, + av_idx, + uc_k_buf, + uc_v_buf, + qk_buf, + softmax_buf, + av_buf, + fd_resources](WebGPUGraph& gr) { + const SdpaLiveState state = compute_live_state(gr); + + wgpuQueueWriteBuffer( + gr.queue(), + uc_k_buf, + 0, + &state.update_cache, + sizeof(state.update_cache)); + wgpuQueueWriteBuffer( + gr.queue(), + uc_v_buf, + 0, + &state.update_cache, + sizeof(state.update_cache)); + if (record_materialized) { + wgpuQueueWriteBuffer( + gr.queue(), + qk_buf, + 0, + &state.attn_weights, + sizeof(state.attn_weights)); + wgpuQueueWriteBuffer( + gr.queue(), softmax_buf, 0, &state.softmax, sizeof(state.softmax)); + wgpuQueueWriteBuffer( + gr.queue(), av_buf, 0, &state.compute_out, sizeof(state.compute_out)); } - const uint32_t kv_off = static_cast( - static_cast(pos) * static_cast(Hkv) * - static_cast(D)); - const uint64_t aw_floats = static_cast(Hq) * - static_cast(s) * static_cast(ctx); - if (aw_floats > UINT32_MAX) { - throw std::runtime_error("WebGPU sdpa: Hq*S*context_len exceeds uint32"); + if (record_fd) { + write_sdpa_fd_decode_uniforms(gr.queue(), fd_resources, state.fd); } - const uint64_t kv_numel = static_cast(s) * - static_cast(Hkv) * static_cast(D); - if (kv_numel > UINT32_MAX) { - throw std::runtime_error("WebGPU sdpa: S*Hkv*D exceeds uint32"); - } - const uint64_t k_cache_numel = static_cast(Cmax) * - static_cast(Hkv) * static_cast(D); - // update_cache K/V: dispatch (kv_numel) + dst offset scale with live S/pos. - UpdateCacheParams uc = - make_update_cache_params(kv_numel, kv_off, k_cache_numel); - wgpuQueueWriteBuffer(gr.queue(), uc_k_buf, 0, &uc, sizeof(uc)); - wgpuQueueWriteBuffer(gr.queue(), uc_v_buf, 0, &uc, sizeof(uc)); - const uint32_t uc_wgc = utils::compute_1d_workgroup_count( - gr.device(), static_cast(kv_numel), uc_wg, "uc(resize)"); - gr.dispatch_at(uc_k_idx).workgroup_count_x = uc_wgc; - gr.dispatch_at(uc_v_idx).workgroup_count_x = uc_wgc; - - // QK: one thread per TM x TN tile; grid = Hq*ceil(S/TM)*ceil(ctx/TN). - AttnWeightsParams qp = - make_attn_weights_params(s, Hq, Hkv, D, ctx, pos, g, scale); - wgpuQueueWriteBuffer(gr.queue(), qk_buf, 0, &qp, sizeof(qp)); - const int64_t qk_tiles = - Hq * utils::div_up(s, kSdpaTileM) * utils::div_up(ctx, kSdpaTileN); - const utils::WgCount qk_wgc = utils::compute_2d_workgroup_count( - gr.device(), static_cast(qk_tiles), qk_wg, "QK(resize)"); - gr.dispatch_at(qk_idx).workgroup_count_x = qk_wgc.x; - gr.dispatch_at(qk_idx).workgroup_count_y = qk_wgc.y; - - // softmax: one workgroup per (h,s) row. - SoftmaxParams sp = make_softmax_params(Hq, s, ctx); - wgpuQueueWriteBuffer(gr.queue(), softmax_buf, 0, &sp, sizeof(sp)); - const utils::WgCount sm_wgc = utils::compute_2d_workgroup_count( - gr.device(), static_cast(Hq * s), 1, "softmax(resize)"); - gr.dispatch_at(softmax_idx).workgroup_count_x = sm_wgc.x; - gr.dispatch_at(softmax_idx).workgroup_count_y = sm_wgc.y; - - // AV: one thread per TM x TN tile; grid = Hq*ceil(S/TM)*ceil(D/TN). - ComputeOutParams op = make_compute_out_params(s, Hq, Hkv, D, ctx, g); - wgpuQueueWriteBuffer(gr.queue(), av_buf, 0, &op, sizeof(op)); - const int64_t av_tiles = - Hq * utils::div_up(s, kSdpaTileM) * utils::div_up(D, kSdpaTileN); - const utils::WgCount av_wgc = utils::compute_2d_workgroup_count( - gr.device(), static_cast(av_tiles), av_wg, "AV(resize)"); - gr.dispatch_at(av_idx).workgroup_count_x = av_wgc.x; - gr.dispatch_at(av_idx).workgroup_count_y = av_wgc.y; - - // Output attn has the same shape as q: [.., S, Hq, D]. + gr.dispatch_at(uc_k_idx).workgroup_count_x = state.update_cache_grid.x; + gr.dispatch_at(uc_k_idx).workgroup_count_y = state.update_cache_grid.y; + gr.dispatch_at(uc_v_idx).workgroup_count_x = state.update_cache_grid.x; + gr.dispatch_at(uc_v_idx).workgroup_count_y = state.update_cache_grid.y; + if (dual_route) { + const size_t active_route = state.use_fd ? 1 : 0; + 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}; + gr.select_dispatch_route(route_group, active_route, active_grids); + } else if (state.use_fd) { + if (!fixed_use_fd) { + throw std::runtime_error("WebGPU sdpa: static route changed"); + } + gr.dispatch_at(fd_resources.dispatch_range.begin).workgroup_count_x = + state.fd.split_grid.x; + gr.dispatch_at(fd_resources.dispatch_range.begin).workgroup_count_y = + state.fd.split_grid.y; + gr.dispatch_at(fd_resources.dispatch_range.begin + 1).workgroup_count_x = + state.fd.reduce_grid.x; + gr.dispatch_at(fd_resources.dispatch_range.begin + 1).workgroup_count_y = + state.fd.reduce_grid.y; + } else { + if (fixed_use_fd) { + throw std::runtime_error("WebGPU sdpa: static route changed"); + } + gr.dispatch_at(qk_idx).workgroup_count_x = state.qk_grid.x; + gr.dispatch_at(qk_idx).workgroup_count_y = state.qk_grid.y; + gr.dispatch_at(softmax_idx).workgroup_count_x = state.softmax_grid.x; + gr.dispatch_at(softmax_idx).workgroup_count_y = state.softmax_grid.y; + gr.dispatch_at(av_idx).workgroup_count_x = state.av_grid.x; + gr.dispatch_at(av_idx).workgroup_count_y = state.av_grid.y; + } gr.set_cur_dims(out_id, gr.cur_dims(q_id)); }; - // q and input_pos share one idempotent recompute; a double-fire is harmless. - graph.add_tensor_resize_hook(q_id, sdpa_resize); + + refresh_state(graph); + graph.add_tensor_resize_hook(q_id, refresh_state); if (dynamic_pos) { - graph.add_resize_hook(input_pos_id, sdpa_resize); + graph.add_resize_hook(input_pos_id, refresh_state); } } diff --git a/backends/webgpu/runtime/ops/sdpa_fd_decode/SdpaFdDecode.cpp b/backends/webgpu/runtime/ops/sdpa_fd_decode/SdpaFdDecode.cpp index ffd3b24dce3..4fa93b0d9e6 100644 --- a/backends/webgpu/runtime/ops/sdpa_fd_decode/SdpaFdDecode.cpp +++ b/backends/webgpu/runtime/ops/sdpa_fd_decode/SdpaFdDecode.cpp @@ -17,6 +17,7 @@ #include +#include #include #include #include @@ -69,6 +70,7 @@ void build_dispatch( WGPUBuffer uniform_buffer, uint64_t uniform_size, uint32_t workgroup_count_x, + bool retain_uniform, const char* kernel_name) { WGPUDevice device = graph.device(); @@ -136,26 +138,55 @@ void build_dispatch( wgpuShaderModuleRelease(shader); wgpuBindGroupLayoutRelease(bgl); wgpuPipelineLayoutRelease(pipeline_layout); - wgpuBufferRelease(uniform_buffer); + if (retain_uniform) { + graph.own_uniform_buffer(uniform_buffer); + } else { + wgpuBufferRelease(uniform_buffer); + } +} + +FdSplitParams make_split_params(const SdpaFdDecodeState& state) { + FdSplitParams params = {}; + params.Hkv = state.Hkv; + params.D = state.D; + params.context_len = state.context_len; + params.g = state.g; + params.num_splits = state.num_splits; + params.split_len = state.split_len; + params.scale = state.scale; + return params; +} + +FdReduceParams make_reduce_params(const SdpaFdDecodeState& state) { + FdReduceParams params = {}; + params.D = state.D; + params.num_splits = state.num_splits; + return params; } } // namespace -void sdpa_fd_decode_dispatch( - WebGPUGraph& graph, - const WebGPUTensor& q, - const WebGPUTensor& k_cache, - const WebGPUTensor& v_cache, - const WebGPUTensor& out, +SdpaFdDecodeState make_sdpa_fd_decode_state( + WGPUDevice device, int64_t Hq, int64_t Hkv, int64_t D, int64_t context_len, int64_t g, float scale) { - // Defensive contract guard: the Sdpa.cpp gate only routes D <= this here, but - // keep the check (lane-owns-D reach) so a future caller can't silently - // overrun. + if (Hq <= 0 || Hkv <= 0 || D <= 0 || context_len <= 0 || g <= 0) { + throw std::runtime_error( + "WebGPU sdpa FlashDecoding: dimensions must be positive"); + } + if (Hq != Hkv * g) { + throw std::runtime_error( + "WebGPU sdpa FlashDecoding: inconsistent GQA dimensions"); + } + if (Hq > UINT32_MAX || Hkv > UINT32_MAX || D > UINT32_MAX || + context_len > UINT32_MAX || g > UINT32_MAX) { + throw std::runtime_error( + "WebGPU sdpa FlashDecoding: parameter exceeds uint32 max"); + } if (D > kSdpaFdMaxHeadDim) { throw std::runtime_error( "WebGPU sdpa FlashDecoding: head dim must be <= " + @@ -165,26 +196,59 @@ void sdpa_fd_decode_dispatch( throw std::runtime_error( "WebGPU sdpa FlashDecoding: head dim must be a multiple of 4"); } - // context_len 0 -> split_len 0 -> empty KV loop -> silent zero output; the - // Sdpa.cpp gate guarantees ctx >= 1, but fail loud if called directly. - if (context_len <= 0) { - throw std::runtime_error( - "WebGPU sdpa FlashDecoding: context_len must be positive"); - } - // Split factor: one split per kSdpaFdSplitTile KV rows, capped. uint32_t num_splits = static_cast( (context_len + kSdpaFdSplitTile - 1) / kSdpaFdSplitTile); - if (num_splits > kSdpaFdMaxSplits) { - num_splits = kSdpaFdMaxSplits; - } + num_splits = std::min(num_splits, kSdpaFdMaxSplits); const uint32_t split_len = static_cast((context_len + num_splits - 1) / num_splits); + const uint64_t split_threads = static_cast(Hq) * + static_cast(num_splits) * + static_cast(kSdpaFdSplitWorkgroupSizeX); + const uint64_t reduce_threads = + static_cast(Hq) * kSdpaFdReduceWorkgroupSizeX; + if (split_threads > UINT32_MAX || reduce_threads > UINT32_MAX) { + throw std::runtime_error( + "WebGPU sdpa FlashDecoding: thread count exceeds uint32 max"); + } + + const uint32_t split_wgc = utils::compute_1d_workgroup_count( + device, + static_cast(split_threads), + kSdpaFdSplitWorkgroupSizeX, + "fd_split"); + const uint32_t reduce_wgc = utils::compute_1d_workgroup_count( + device, + static_cast(reduce_threads), + kSdpaFdReduceWorkgroupSizeX, + "fd_reduce"); + return { + static_cast(Hq), + static_cast(Hkv), + static_cast(D), + static_cast(context_len), + static_cast(g), + num_splits, + split_len, + scale, + {split_wgc, 1u}, + {reduce_wgc, 1u}}; +} + +SdpaFdDecodeResources record_sdpa_fd_decode_dispatches( + WebGPUGraph& graph, + const WebGPUTensor& q, + const WebGPUTensor& k_cache, + const WebGPUTensor& v_cache, + const WebGPUTensor& out, + const SdpaFdDecodeState& state) { + const size_t dispatch_begin = graph.num_dispatches(); + // Scratch: per-(head,split) partials at kSdpaFdMaxSplits stride. - const uint64_t po_floats = static_cast(Hq) * - static_cast(kSdpaFdMaxSplits) * static_cast(D); - const uint64_t pml_floats = static_cast(Hq) * + const uint64_t po_floats = static_cast(state.Hq) * + static_cast(kSdpaFdMaxSplits) * static_cast(state.D); + const uint64_t pml_floats = static_cast(state.Hq) * static_cast(kSdpaFdMaxSplits) * 2ull; WGPUBuffer part_o = graph.acquire_scratch(po_floats * sizeof(float)); WebGPUGraph::ScopedScratch part_o_guard(&graph, part_o); @@ -192,14 +256,7 @@ void sdpa_fd_decode_dispatch( WebGPUGraph::ScopedScratch part_ml_guard(&graph, part_ml); // Pass 1: split (Hq*num_splits WGs) -> writes part_o, part_ml. - FdSplitParams sp = {}; - sp.Hkv = static_cast(Hkv); - sp.D = static_cast(D); - sp.context_len = static_cast(context_len); - sp.g = static_cast(g); - sp.num_splits = num_splits; - sp.split_len = split_len; - sp.scale = scale; + FdSplitParams sp = make_split_params(state); WGPUBuffer ub_split = graph.make_uniform_buffer(&sp, sizeof(sp)); BufferBinding split_bindings[5] = { {part_o, po_floats * sizeof(float)}, @@ -207,20 +264,6 @@ void sdpa_fd_decode_dispatch( {q.buffer, q.nbytes}, {k_cache.buffer, k_cache.nbytes}, {v_cache.buffer, v_cache.nbytes}}; - // Compute the thread product in 64-bit + guard before the u32 cast, mirroring - // the Sdpa.cpp aw_floats > UINT32_MAX guards. - const uint64_t split_threads = static_cast(Hq) * - static_cast(num_splits) * - static_cast(kSdpaFdSplitWorkgroupSizeX); - if (split_threads > UINT32_MAX) { - throw std::runtime_error( - "WebGPU sdpa FlashDecoding: split thread count exceeds uint32 max"); - } - const uint32_t wgc_split = utils::compute_1d_workgroup_count( - graph.device(), - static_cast(split_threads), - kSdpaFdSplitWorkgroupSizeX, - "fd_split"); const char* split_shader = kSdpaFdSplitWGSL; if (graph.kv_f16()) { split_shader = kSdpaFdSplitHalfWGSL; @@ -233,23 +276,17 @@ void sdpa_fd_decode_dispatch( 2, ub_split, sizeof(sp), - wgc_split, + state.split_grid.x, + true, "fd_split"); // Pass 2: reduce (Hq WGs) -> reads part_o, part_ml; writes out. - FdReduceParams rp = {}; - rp.D = static_cast(D); - rp.num_splits = num_splits; + FdReduceParams rp = make_reduce_params(state); WGPUBuffer ub_reduce = graph.make_uniform_buffer(&rp, sizeof(rp)); BufferBinding reduce_bindings[3] = { {out.buffer, out.nbytes}, {part_o, po_floats * sizeof(float)}, {part_ml, pml_floats * sizeof(float)}}; - const uint32_t wgc_reduce = utils::compute_1d_workgroup_count( - graph.device(), - static_cast(Hq) * kSdpaFdReduceWorkgroupSizeX, - kSdpaFdReduceWorkgroupSizeX, - "fd_reduce"); build_dispatch( graph, kSdpaFdReduceWGSL, @@ -258,8 +295,27 @@ void sdpa_fd_decode_dispatch( 1, ub_reduce, sizeof(rp), - wgc_reduce, + state.reduce_grid.x, + true, "fd_reduce"); + + return {ub_split, ub_reduce, {dispatch_begin, graph.num_dispatches()}}; +} + +void write_sdpa_fd_decode_uniforms( + WGPUQueue queue, + const SdpaFdDecodeResources& resources, + const SdpaFdDecodeState& state) { + FdSplitParams split_params = make_split_params(state); + FdReduceParams reduce_params = make_reduce_params(state); + wgpuQueueWriteBuffer( + queue, resources.split_uniform, 0, &split_params, sizeof(split_params)); + wgpuQueueWriteBuffer( + queue, + resources.reduce_uniform, + 0, + &reduce_params, + sizeof(reduce_params)); } } // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/sdpa_fd_decode/SdpaFdDecode.h b/backends/webgpu/runtime/ops/sdpa_fd_decode/SdpaFdDecode.h index 1b161cd8ec0..d86b4c624b8 100644 --- a/backends/webgpu/runtime/ops/sdpa_fd_decode/SdpaFdDecode.h +++ b/backends/webgpu/runtime/ops/sdpa_fd_decode/SdpaFdDecode.h @@ -9,6 +9,7 @@ #pragma once #include +#include #include @@ -19,15 +20,27 @@ namespace executorch::backends::webgpu { // SDPA path (the FD selection predicate in Sdpa.cpp checks this). constexpr int64_t kSdpaFdMaxHeadDim = 128; -// Split-KV FlashDecoding decode dispatch (S==1): a split pass over -// Hq*num_splits workgroups + a reduce pass over Hq workgroups. Called from the -// Sdpa.cpp WEBGPU_SDPA_FD branch. -void sdpa_fd_decode_dispatch( - WebGPUGraph& graph, - const WebGPUTensor& q, - const WebGPUTensor& k_cache, - const WebGPUTensor& v_cache, - const WebGPUTensor& out, +struct SdpaFdDecodeState { + uint32_t Hq; + uint32_t Hkv; + uint32_t D; + uint32_t context_len; + uint32_t g; + uint32_t num_splits; + uint32_t split_len; + float scale; + utils::WgCount split_grid; + utils::WgCount reduce_grid; +}; + +struct SdpaFdDecodeResources { + WGPUBuffer split_uniform; + WGPUBuffer reduce_uniform; + utils::DispatchRange dispatch_range; +}; + +SdpaFdDecodeState make_sdpa_fd_decode_state( + WGPUDevice device, int64_t Hq, int64_t Hkv, int64_t D, @@ -35,4 +48,19 @@ void sdpa_fd_decode_dispatch( int64_t g, float scale); +// Records split + reduce with retained UBOs. Route selection is owned by the +// caller so this helper never mutates recorded dispatch counts. +SdpaFdDecodeResources record_sdpa_fd_decode_dispatches( + WebGPUGraph& graph, + const WebGPUTensor& q, + const WebGPUTensor& k_cache, + const WebGPUTensor& v_cache, + const WebGPUTensor& out, + const SdpaFdDecodeState& state); + +void write_sdpa_fd_decode_uniforms( + WGPUQueue queue, + const SdpaFdDecodeResources& resources, + const SdpaFdDecodeState& state); + } // namespace executorch::backends::webgpu diff --git a/backends/webgpu/test/native/test_dispatch_2d.cpp b/backends/webgpu/test/native/test_dispatch_2d.cpp index 582caa3c98e..305a6c39ee0 100644 --- a/backends/webgpu/test/native/test_dispatch_2d.cpp +++ b/backends/webgpu/test/native/test_dispatch_2d.cpp @@ -9,6 +9,7 @@ // Device-free unit test for the pure 2D workgroup-count fold that lifts the // 65535 per-dim dispatch cap. Exercises the fold arithmetic only — no GPU. +#include #include #include @@ -16,6 +17,10 @@ #include #include +using executorch::backends::webgpu::WebGPUDispatch; +using executorch::backends::webgpu::WebGPUGraph; +using executorch::backends::webgpu::utils::DispatchRange; +using executorch::backends::webgpu::utils::DispatchRouteRegistry; using executorch::backends::webgpu::utils::fold_workgroup_count_2d; using executorch::backends::webgpu::utils::WgCount; @@ -57,4 +62,157 @@ TEST(DispatchFold, ThrowsWhenNeeds3rdDimension) { EXPECT_ANY_THROW(fold_workgroup_count_2d(kMax * kMax + 1u, kMax, "test")); } +void expect_grid(const WgCount& grid, uint32_t x, uint32_t y) { + EXPECT_EQ(grid.x, x); + EXPECT_EQ(grid.y, y); +} + +void expect_grids_equal( + const std::vector& actual, + const std::vector& expected) { + ASSERT_EQ(actual.size(), expected.size()); + for (size_t i = 0; i < actual.size(); i++) { + EXPECT_EQ(actual[i].x, expected[i].x) << "index=" << i; + EXPECT_EQ(actual[i].y, expected[i].y) << "index=" << i; + } +} + +TEST(DispatchRoute, SwitchesUnequalRangesAndRestoresBothDimensions) { + std::vector grids(6, {101, 103}); + const std::vector compute(6, true); + DispatchRouteRegistry registry; + const size_t group = registry.register_group( + grids.size(), {{1, 3}, {4, 5}}, [&](size_t i) { return compute[i]; }); + auto set_grid = [&](size_t i, WgCount grid) { grids[i] = grid; }; + + registry.select(group, 0, {{3, 7}, {5, 11}}, set_grid); + expect_grid(grids[0], 101, 103); + expect_grid(grids[1], 3, 7); + expect_grid(grids[2], 5, 11); + expect_grid(grids[3], 101, 103); + expect_grid(grids[4], 0, 0); + expect_grid(grids[5], 101, 103); + + registry.select(group, 1, {{13, 17}}, set_grid); + expect_grid(grids[1], 0, 0); + expect_grid(grids[2], 0, 0); + expect_grid(grids[4], 13, 17); + + registry.select(group, 0, {{19, 23}, {29, 31}}, set_grid); + expect_grid(grids[1], 19, 23); + expect_grid(grids[2], 29, 31); + expect_grid(grids[4], 0, 0); +} + +TEST(DispatchRoute, InvalidSelectionDoesNotMutateGrids) { + std::vector grids = {{1, 2}, {3, 4}, {5, 6}}; + const std::vector original = grids; + DispatchRouteRegistry registry; + const size_t group = registry.register_group( + grids.size(), {{0, 2}, {2, 3}}, [](size_t) { return true; }); + auto set_grid = [&](size_t i, WgCount grid) { grids[i] = grid; }; + + EXPECT_ANY_THROW(registry.select(group, 2, {{7, 8}}, set_grid)); + expect_grids_equal(grids, original); + EXPECT_ANY_THROW(registry.select(group, 0, {{7, 8}}, set_grid)); + expect_grids_equal(grids, original); + EXPECT_ANY_THROW(registry.select(group, 0, {{0, 8}, {9, 10}}, set_grid)); + expect_grids_equal(grids, original); + EXPECT_ANY_THROW(registry.select(group, 0, {{7, 0}, {9, 10}}, set_grid)); + expect_grids_equal(grids, original); +} + +TEST(DispatchRoute, RejectsInvalidAndMultiplyOwnedRanges) { + const auto all_compute = [](size_t) { return true; }; + DispatchRouteRegistry registry; + EXPECT_ANY_THROW(registry.register_group(4, {{2, 1}}, all_compute)); + EXPECT_ANY_THROW(registry.register_group(4, {{1, 1}}, all_compute)); + EXPECT_ANY_THROW(registry.register_group(4, {{0, 5}}, all_compute)); + EXPECT_ANY_THROW(registry.register_group(4, {{0, 2}, {1, 3}}, all_compute)); + + const size_t first = + registry.register_group(4, {{0, 1}, {1, 2}}, all_compute); + EXPECT_EQ(first, 0); + EXPECT_ANY_THROW(registry.register_group(4, {{1, 3}}, all_compute)); + + DispatchRouteRegistry copy_registry; + const std::vector compute = {true, false, true}; + EXPECT_ANY_THROW(copy_registry.register_group( + compute.size(), {{0, 1}, {1, 2}}, [&](size_t i) { return compute[i]; })); +} + +TEST(DispatchRoute, GraphRejectsCopyAndCrossGroupOwnership) { + WebGPUGraph graph; + for (size_t i = 0; i < 6; i++) { + graph.add_dispatch(WebGPUDispatch{}); + } + graph.add_buffer_copy(nullptr, nullptr, 0); + + const size_t group = graph.register_dispatch_route_group({{0, 1}, {1, 2}}); + EXPECT_EQ(group, 0); + EXPECT_ANY_THROW(graph.register_dispatch_route_group({{1, 2}, {2, 3}})); + EXPECT_ANY_THROW(graph.register_dispatch_route_group({{5, 6}, {6, 7}})); + + const size_t second = graph.register_dispatch_route_group({{2, 4}, {4, 5}}); + EXPECT_EQ(second, 1); + + graph.select_dispatch_route(group, 0, {{7, 11}}); + graph.select_dispatch_route(second, 0, {{13, 17}, {19, 23}}); + expect_grid( + {graph.dispatch_at(0).workgroup_count_x, + graph.dispatch_at(0).workgroup_count_y}, + 7, + 11); + expect_grid( + {graph.dispatch_at(1).workgroup_count_x, + graph.dispatch_at(1).workgroup_count_y}, + 0, + 0); + expect_grid( + {graph.dispatch_at(2).workgroup_count_x, + graph.dispatch_at(2).workgroup_count_y}, + 13, + 17); + expect_grid( + {graph.dispatch_at(3).workgroup_count_x, + graph.dispatch_at(3).workgroup_count_y}, + 19, + 23); + expect_grid( + {graph.dispatch_at(4).workgroup_count_x, + graph.dispatch_at(4).workgroup_count_y}, + 0, + 0); + + graph.select_dispatch_route(group, 1, {{29, 31}}); + expect_grid( + {graph.dispatch_at(2).workgroup_count_x, + graph.dispatch_at(2).workgroup_count_y}, + 13, + 17); +} + +TEST(DispatchRoute, ExecuteRejectsHalfZeroGrid) { + WebGPUGraph graph; + WebGPUDispatch dispatch; + dispatch.workgroup_count_x = 0; + dispatch.workgroup_count_y = 1; + graph.add_dispatch(dispatch); + EXPECT_ANY_THROW(graph.execute({})); +} + +TEST(DispatchRoute, RecordsAlternatesOnlyForDynamicEligibleGraphs) { + using executorch::backends::webgpu::utils::should_record_q4gsw_dual_route; + using executorch::backends::webgpu::utils::should_record_sdpa_dual_route; + + EXPECT_FALSE(should_record_q4gsw_dual_route(32, true, false)); + EXPECT_FALSE(should_record_q4gsw_dual_route(1, true, true)); + EXPECT_FALSE(should_record_q4gsw_dual_route(32, false, true)); + EXPECT_TRUE(should_record_q4gsw_dual_route(32, true, 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)); +} + } // namespace diff --git a/backends/webgpu/test/native/test_dynamic_shape.cpp b/backends/webgpu/test/native/test_dynamic_shape.cpp index f41adbf09ec..874feda5880 100644 --- a/backends/webgpu/test/native/test_dynamic_shape.cpp +++ b/backends/webgpu/test/native/test_dynamic_shape.cpp @@ -24,11 +24,13 @@ // /tmp/dynamic_shape. #include +#include #include #include #include +#include #include #include #include @@ -49,6 +51,26 @@ constexpr int kHidden = 64; // Artifacts directory; set from env/argv in main() before RUN_ALL_TESTS(). std::string g_dir; // NOLINT(cppcoreguidelines-avoid-non-const-global-variables) +#ifdef WGPU_BACKEND_ENABLE_PROFILING +std::vector current_profile_names() { + const auto* context = get_default_webgpu_context(); + if (context == nullptr || context->querypool == nullptr) { + return {}; + } + std::vector names; + for (const auto& duration : context->querypool->results()) { + names.push_back(duration.kernel_name); + } + return names; +} + +bool contains_name( + const std::vector& names, + const std::string& expected) { + return std::find(names.begin(), names.end(), expected) != names.end(); +} +#endif + std::vector read_bin(const std::string& path) { std::ifstream f(path, std::ios::binary | std::ios::ate); if (!f) { @@ -105,16 +127,22 @@ void check_s(Module& m, const std::string& prefix, int s) { // 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; +constexpr int kLinAltK = 72; constexpr int kLinN = 128; constexpr int kLinNShmem = 2048; // 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(Module& m, int m_rows, const char* prefix, int n) { +void run_linear( + Module& m, + int m_rows, + const char* prefix, + int n, + int k = kLinK) { const std::string base = g_dir + "/" + prefix + ".S" + std::to_string(m_rows); auto input = read_bin(base + ".input.bin"); auto golden = read_bin(base + ".golden.bin"); ASSERT_FALSE(input.empty()) << "missing " << prefix << ".S" << m_rows; - auto t = make_tensor_ptr({m_rows, kLinK}, std::move(input)); + auto t = make_tensor_ptr({m_rows, k}, std::move(input)); auto r = m.forward({EValue(t)}); ASSERT_TRUE(r.ok() && !r.get().empty() && r.get()[0].isTensor()) << prefix << " M=" << m_rows << " forward failed"; @@ -138,16 +166,27 @@ void check_linear(int m_rows) { void check_linear_shmem(int m_rows) { Module m(g_dir + "/dyn_linear_shmem.pte"); ASSERT_EQ(m.load_forward(), Error::Ok) << "load dyn_linear_shmem.pte"; - run_linear(m, m_rows, "dyn_linear_shmem", kLinNShmem); + run_linear(m, m_rows, "dyn_linear_shmem", kLinNShmem, kLinAltK); +} + +void check_linear_tiled(int m_rows) { + Module m(g_dir + "/dyn_linear_tiled.pte"); + ASSERT_EQ(m.load_forward(), Error::Ok) << "load dyn_linear_tiled.pte"; + run_linear(m, m_rows, "dyn_linear_tiled", kLinN, kLinAltK); } // Dynamic SDPA (GQA prefill, input_pos=0): q[1,s,hq,d] k/v[1,s,hkv,d] // caches[1,cmax,hkv,d]; attn output [1,s,hq,d] selected by shape (3 outputs). constexpr int kSdHq = 8, kSdHkv = 2, kSdD = 16, kSdCmax = 64; -void check_sdpa(int s) { - Module m(g_dir + "/sdpa_dyn.pte"); - ASSERT_EQ(m.load_forward(), Error::Ok) << "sdpa_dyn S=" << s << " load"; - const std::string b = g_dir + "/sdpa_dyn.S" + std::to_string(s) + "."; +void run_sdpa_case( + Module& m, + int s, + const char* prefix, + int hq, + int hkv, + int d, + int cmax) { + 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"); auto v = read_bin(b + "v.bin"); @@ -158,36 +197,92 @@ void check_sdpa(int s) { q.empty() || k.empty() || v.empty() || kc.empty() || vc.empty() || golden.empty()) << "missing sdpa_dyn.S" << s; - auto tq = make_tensor_ptr({1, s, kSdHq, kSdD}, std::move(q)); - auto tk = make_tensor_ptr({1, s, kSdHkv, kSdD}, std::move(k)); - auto tv = make_tensor_ptr({1, s, kSdHkv, kSdD}, std::move(v)); - auto tkc = make_tensor_ptr({1, kSdCmax, kSdHkv, kSdD}, std::move(kc)); - auto tvc = make_tensor_ptr({1, kSdCmax, kSdHkv, kSdD}, std::move(vc)); + 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)); + auto tkc = make_tensor_ptr({1, cmax, hkv, d}, std::move(kc)); + auto tvc = make_tensor_ptr({1, cmax, hkv, d}, std::move(vc)); auto r = m.forward({EValue(tq), EValue(tk), EValue(tv), EValue(tkc), EValue(tvc)}); ASSERT_TRUE(r.ok()) << "sdpa S=" << s << " forward failed (err=" << (int)r.error() << ")"; // Select the attn output by full shape [1,s,hq,d] (never numel). const float* attn = nullptr; - const size_t numel = static_cast(s) * kSdHq * kSdD; + const size_t numel = static_cast(s) * hq * d; for (size_t i = 0; i < r.get().size(); i++) { if (!r.get()[i].isTensor()) { continue; } const auto& t = r.get()[i].toTensor(); - if (t.dim() == 4 && t.size(1) == s && t.size(2) == kSdHq && - t.size(3) == kSdD) { + if (t.dim() == 4 && t.size(1) == s && t.size(2) == hq && t.size(3) == d) { attn = t.const_data_ptr(); break; } } ASSERT_NE(attn, nullptr) << "sdpa S=" << s << ": no attn output of shape [1," - << s << "," << kSdHq << "," << kSdD << "]"; + << 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; } +void run_sdpa(Module& m, int s) { + run_sdpa_case(m, s, "sdpa_dyn", kSdHq, kSdHkv, kSdD, kSdCmax); +} + +void check_sdpa(int s) { + Module m(g_dir + "/sdpa_dyn.pte"); + ASSERT_EQ(m.load_forward(), Error::Ok) << "sdpa_dyn S=" << s << " load"; + run_sdpa(m, s); +} + +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"); + auto q = read_bin(b + "q.bin"); + auto k = read_bin(b + "k.bin"); + auto v = read_bin(b + "v.bin"); + auto kc = read_bin(b + "kc.bin"); + auto vc = read_bin(b + "vc.bin"); + auto golden = read_bin(b + "golden.bin"); + ASSERT_FALSE( + x.empty() || q.empty() || k.empty() || v.empty() || kc.empty() || + vc.empty() || golden.empty()) + << "missing combined_routes.S" << s; + auto tx = make_tensor_ptr({s, kLinK}, std::move(x)); + auto tq = make_tensor_ptr({1, s, kSdHq, kSdD}, std::move(q)); + auto tk = make_tensor_ptr({1, s, kSdHkv, kSdD}, std::move(k)); + auto tv = make_tensor_ptr({1, s, kSdHkv, kSdD}, std::move(v)); + auto tkc = make_tensor_ptr({1, kSdCmax, kSdHkv, kSdD}, std::move(kc)); + auto tvc = make_tensor_ptr({1, kSdCmax, kSdHkv, kSdD}, std::move(vc)); + auto result = m.forward( + {EValue(tx), + EValue(tq), + EValue(tk), + EValue(tv), + EValue(tkc), + EValue(tvc)}); + ASSERT_TRUE(result.ok()) << "combined routes S=" << s + << " forward failed (err=" << (int)result.error() + << ")"; + const float* attn = nullptr; + const size_t numel = static_cast(s) * kSdHq * kSdD; + for (const auto& output : result.get()) { + if (!output.isTensor()) { + continue; + } + const auto& tensor = output.toTensor(); + if (tensor.dim() == 4 && tensor.size(1) == s && tensor.size(2) == kSdHq && + tensor.size(3) == kSdD) { + attn = tensor.const_data_ptr(); + break; + } + } + ASSERT_NE(attn, nullptr); + const std::vector got(attn, attn + numel); + EXPECT_LT(max_err(got, golden), 1e-2f) << "combined_routes S=" << s; +} + // Dynamic embedding: int64 token ids [N] -> [N, kEmbDim] fp32. The int64 host // input exercises copy_inputs' int64->int32 narrow path under dynamic shapes. constexpr int kEmbDim = 64; @@ -407,9 +502,7 @@ TEST(DynamicShape, QuantizedLinearReusedGraph) { } } -// I3: dynamic linear at N=2048 -> the shmem-GEMM route (K>=4096||N>=2048); the -// resize hook recomputes the shmem tile count for the live M on the fixed shmem -// pipeline (M=1 exercises a partial row-tile). +// I3: K=72 disables Steel; N=2048 forces shmem for M>1 and bicol for M=1. TEST(DynamicShape, QuantizedLinearShmem) { for (int m_rows : {128, 32, 1}) { check_linear_shmem(m_rows); @@ -421,26 +514,172 @@ TEST(DynamicShape, QuantizedLinearShmemReusedGraph) { Module m(g_dir + "/dyn_linear_shmem.pte"); ASSERT_EQ(m.load_forward(), Error::Ok) << "load dyn_linear_shmem.pte"; for (int m_rows : {128, 32, 1, 128}) { - run_linear(m, m_rows, "dyn_linear_shmem", kLinNShmem); + run_linear(m, m_rows, "dyn_linear_shmem", kLinNShmem, kLinAltK); } } -// J: dynamic SDPA (GQA prefill) at several seq-len S. The whole case skips -// while op coverage is pending (the dynamic-S build throws err 48 until -// registered). -TEST(DynamicShape, Sdpa) { - { - Module probe(g_dir + "/sdpa_dyn.pte"); - if (probe.load_forward() == Error::DelegateInvalidCompatibility) { - GTEST_SKIP() << "sdpa_dyn pending op coverage (err " - << (int)Error::DelegateInvalidCompatibility << ")"; - } +// I5: K=72 disables Steel; N=128 keeps the tiled fallback for M>1. +TEST(DynamicShape, QuantizedLinearTiled) { + for (int m_rows : {128, 32, 1}) { + check_linear_tiled(m_rows); + } +} + +TEST(DynamicShape, QuantizedLinearTiledReusedGraph) { + Module m(g_dir + "/dyn_linear_tiled.pte"); + ASSERT_EQ(m.load_forward(), Error::Ok) << "load dyn_linear_tiled.pte"; + for (int m_rows : {128, 1, 32, 1, 128}) { + run_linear(m, m_rows, "dyn_linear_tiled", kLinN, kLinAltK); + } +} + +#ifdef WGPU_BACKEND_ENABLE_PROFILING +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"; + } + Module m(g_dir + "/combined_routes.pte"); + ASSERT_EQ(m.load_forward(), Error::Ok) << "load combined_routes.pte"; + for (int s : {64, 1, 16, 1, 64}) { + run_combined_routes(m, s); + const auto names = current_profile_names(); + ASSERT_EQ(names.size(), s == 1 ? 6 : 7); + EXPECT_EQ(std::count(names.begin(), names.end(), ""), 1); + EXPECT_EQ( + std::count(names.begin(), names.end(), "linear_q4gsw_coop4_bicol"), + s == 1 ? 1 : 0); + EXPECT_EQ( + std::count(names.begin(), names.end(), "linear_q4gsw_steel"), + s != 1 ? 1 : 0); + EXPECT_FALSE(contains_name(names, "linear_q4gsw_shmem")); + EXPECT_FALSE(contains_name(names, "linear_q4gsw_tiled")); + EXPECT_EQ(std::count(names.begin(), names.end(), "update_cache"), 2); + EXPECT_EQ(std::count(names.begin(), names.end(), "fd_split"), s == 1); + EXPECT_EQ(std::count(names.begin(), names.end(), "fd_reduce"), s == 1); + EXPECT_EQ( + std::count(names.begin(), names.end(), "sdpa_compute_attn_weights"), + s != 1); + EXPECT_EQ(std::count(names.begin(), names.end(), "sdpa_softmax"), s != 1); + EXPECT_EQ( + std::count(names.begin(), names.end(), "sdpa_compute_out"), s != 1); + } +} + +TEST(DynamicShape, StaticRouteProfiles) { + const auto* context = get_default_webgpu_context(); + if (std::getenv("WEBGPU_TIMESTAMP_QUERY") == nullptr || context == nullptr || + !context->timestamp_supported) { + GTEST_SKIP() << "timestamp queries unavailable"; } + + Module linear_m1(g_dir + "/static_linear_m1.pte"); + ASSERT_EQ(linear_m1.load_forward(), Error::Ok); + run_linear(linear_m1, 1, "static_linear_m1", kLinN); + auto names = current_profile_names(); + ASSERT_EQ(names.size(), 1); + EXPECT_EQ( + std::count(names.begin(), names.end(), "linear_q4gsw_coop4_bicol"), 1); + + Module linear_m32(g_dir + "/static_linear_m32.pte"); + ASSERT_EQ(linear_m32.load_forward(), Error::Ok); + run_linear(linear_m32, 32, "static_linear_m32", kLinN); + names = current_profile_names(); + ASSERT_EQ(names.size(), 1); + EXPECT_EQ(std::count(names.begin(), names.end(), "linear_q4gsw_steel"), 1); + + Module linear_shmem(g_dir + "/dyn_linear_shmem.pte"); + ASSERT_EQ(linear_shmem.load_forward(), Error::Ok); + run_linear(linear_shmem, 128, "dyn_linear_shmem", kLinNShmem, kLinAltK); + names = current_profile_names(); + ASSERT_EQ(names.size(), 1); + EXPECT_EQ(std::count(names.begin(), names.end(), "linear_q4gsw_shmem"), 1); + + Module linear_tiled(g_dir + "/dyn_linear_tiled.pte"); + ASSERT_EQ(linear_tiled.load_forward(), Error::Ok); + run_linear(linear_tiled, 128, "dyn_linear_tiled", kLinN, kLinAltK); + names = current_profile_names(); + ASSERT_EQ(names.size(), 1); + EXPECT_EQ(std::count(names.begin(), names.end(), "linear_q4gsw_tiled"), 1); + + Module sdpa_s1(g_dir + "/static_sdpa_s1.pte"); + ASSERT_EQ(sdpa_s1.load_forward(), Error::Ok); + run_sdpa_case(sdpa_s1, 1, "static_sdpa_s1", kSdHq, kSdHkv, kSdD, kSdCmax); + names = current_profile_names(); + ASSERT_EQ(names.size(), 4); + EXPECT_EQ(std::count(names.begin(), names.end(), "fd_split"), 1); + EXPECT_EQ(std::count(names.begin(), names.end(), "fd_reduce"), 1); + + Module sdpa_s16(g_dir + "/static_sdpa_s16.pte"); + ASSERT_EQ(sdpa_s16.load_forward(), Error::Ok); + run_sdpa_case(sdpa_s16, 16, "static_sdpa_s16", kSdHq, kSdHkv, kSdD, kSdCmax); + names = current_profile_names(); + ASSERT_EQ(names.size(), 5); + EXPECT_EQ( + std::count(names.begin(), names.end(), "sdpa_compute_attn_weights"), 1); + EXPECT_EQ(std::count(names.begin(), names.end(), "sdpa_softmax"), 1); + EXPECT_EQ(std::count(names.begin(), names.end(), "sdpa_compute_out"), 1); +} +#endif + +// J: dynamic SDPA reuses one graph across prefill and FlashDecoding shapes. +TEST(DynamicShape, Sdpa) { for (int s : {64, 16, 1}) { check_sdpa(s); } } +TEST(DynamicShape, SdpaReusedGraph) { + Module m(g_dir + "/sdpa_dyn.pte"); + ASSERT_EQ(m.load_forward(), Error::Ok) << "load sdpa_dyn.pte"; + for (int s : {64, 1, 16, 1, 64}) { + run_sdpa(m, s); + } +} + +TEST(DynamicShape, CombinedLiveRoutes) { + Module m(g_dir + "/combined_routes.pte"); + ASSERT_EQ(m.load_forward(), Error::Ok) << "load combined_routes.pte"; + for (int s : {64, 1, 16, 1, 64}) { + run_combined_routes(m, s); + } +} + +TEST(DynamicShape, SdpaWideMaterializedOnly) { + Module m(g_dir + "/sdpa_wide.pte"); + ASSERT_EQ(m.load_forward(), Error::Ok) << "load sdpa_wide.pte"; + for (int s : {16, 1, 16}) { + run_sdpa_case(m, s, "sdpa_wide", 8, 2, 132, 16); + } +} + +#ifdef WGPU_BACKEND_ENABLE_PROFILING +TEST(DynamicShape, SdpaLiveRoutesProfile) { + const auto* context = get_default_webgpu_context(); + if (std::getenv("WEBGPU_TIMESTAMP_QUERY") == nullptr || context == nullptr || + !context->timestamp_supported) { + GTEST_SKIP() << "timestamp queries unavailable"; + } + Module m(g_dir + "/sdpa_dyn.pte"); + ASSERT_EQ(m.load_forward(), Error::Ok) << "load sdpa_dyn.pte"; + for (int s : {64, 1, 16, 1, 64}) { + run_sdpa(m, s); + const auto names = current_profile_names(); + ASSERT_EQ(names.size(), s == 1 ? 4 : 5); + EXPECT_EQ(std::count(names.begin(), names.end(), "update_cache"), 2); + EXPECT_EQ(std::count(names.begin(), names.end(), "fd_split"), s == 1); + EXPECT_EQ(std::count(names.begin(), names.end(), "fd_reduce"), s == 1); + EXPECT_EQ( + std::count(names.begin(), names.end(), "sdpa_compute_attn_weights"), + s != 1); + EXPECT_EQ(std::count(names.begin(), names.end(), "sdpa_softmax"), s != 1); + EXPECT_EQ( + std::count(names.begin(), names.end(), "sdpa_compute_out"), s != 1); + } +} +#endif + // K: dynamic embedding (int64 token ids) at several token counts. TEST(DynamicShape, Embedding) { for (int n : {16, 8, 1}) { diff --git a/backends/webgpu/test/native/test_execution_options.cpp b/backends/webgpu/test/native/test_execution_options.cpp index 2f96636438d..b927afc5dcf 100644 --- a/backends/webgpu/test/native/test_execution_options.cpp +++ b/backends/webgpu/test/native/test_execution_options.cpp @@ -123,5 +123,28 @@ TEST(WebGPUExecutionPlanTest, RejectsInvalidSuppressibleRange) { std::runtime_error); } +TEST(WebGPUExecutionPlanTest, FiltersDisabledDispatchesAcrossChunks) { + const std::vector enabled = {true, false, true, false, true, true}; + const WebGPUExecutionPlan plan = plan_webgpu_execution( + 6, 1, ExecuteConfig{2, 1}, {}, WebGPUGraphExecutionOptions{}, enabled); + + EXPECT_EQ( + plan.dispatch_chunks, + (std::vector>{{0}, {2}, {4}, {5}})); + EXPECT_EQ(plan.copy_outputs, (std::vector{true})); +} + +TEST(WebGPUExecutionPlanTest, RejectsMismatchedEnabledDispatches) { + EXPECT_THROW( + plan_webgpu_execution( + 3, + 1, + ExecuteConfig{}, + {}, + WebGPUGraphExecutionOptions{}, + {true, false}), + std::runtime_error); +} + } // namespace } // namespace executorch::backends::webgpu 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 969c5e06343..57feaeeb414 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,7 @@ 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 MAXS = 128 # upper bound for the dynamic seq-len dim (within the 1D dispatch cap) HIDDEN = 64 @@ -108,17 +109,29 @@ def _ramp(shape) -> torch.Tensor: return torch.linspace(-1.0, 1.0, n, dtype=torch.float32).reshape(shape) +def _lower_fully_delegated(ep, label: str): + edge = to_edge_transform_and_lower(ep, partitioner=[VulkanPartitioner()]) + graph = edge.exported_program().graph_module.graph + delegates = get_delegates(graph) + portable = get_non_lowered_nodes(graph) + if len(delegates) != 1: + raise RuntimeError(f"{label}: expected one delegate, got {len(delegates)}") + if portable: + raise RuntimeError(f"{label}: non-lowered nodes: {portable}") + 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 RuntimeError(f"{label}: serialized delegates: {delegate_ids}") + return et + + def _export(model, example_inputs, dynamic_shapes, path: str) -> None: ep = torch.export.export(model, example_inputs, dynamic_shapes=dynamic_shapes) - et = to_edge_transform_and_lower( - ep, partitioner=[VulkanPartitioner()] - ).to_executorch() - found = any( - d.id == "VulkanBackend" - for plan in et.executorch_program.execution_plan - for d in plan.delegates - ) - assert found, f"Expected VulkanBackend delegate in {path}" + et = _lower_fully_delegated(ep, path) with open(path, "wb") as f: f.write(et.buffer) print(f"Exported {path}") @@ -186,10 +199,28 @@ def export_dynamic_shape_cases(out_dir: str) -> None: # 2d) 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(out_dir, n=LIN_SHMEM_N, prefix="dyn_linear_shmem") + _export_dynamic_linear( + out_dir, + n=LIN_SHMEM_N, + prefix="dyn_linear_shmem", + k=LIN_ALT_K, + group=LIN_ALT_GROUP, + ) + _export_dynamic_linear( + out_dir, + prefix="dyn_linear_tiled", + k=LIN_ALT_K, + group=LIN_ALT_GROUP, + ) + _export_static_linear(out_dir, 1, "static_linear_m1") + _export_static_linear(out_dir, 32, "static_linear_m32") # 2e) Fused SDPA with a DYNAMIC seq-len S (prefill, input_pos=0). _export_dynamic_sdpa(out_dir) + _export_combined_routes(out_dir) + _export_dynamic_sdpa_wide(out_dir) + _export_static_sdpa(out_dir, 1, "static_sdpa_s1") + _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) @@ -224,36 +255,35 @@ def export_dynamic_shape_cases(out_dir: str) -> None: # Quantized linear: K x N weight, dynamic rows M; input [M, K], output [M, N]. LIN_K = 64 LIN_N = 128 -LIN_SHMEM_N = 2048 # N>=2048 routes linear_q4gsw to the shmem-GEMM path +LIN_ALT_K = 72 # K%16 != 0 disables Steel while K%8 keeps bicol eligible. +LIN_ALT_GROUP = 24 +LIN_SHMEM_N = 2048 LIN_GROUP = 32 LIN_MAXM = 128 def _export_dynamic_linear( - out_dir: str, n: int = LIN_N, prefix: str = "dyn_linear" + out_dir: str, + n: int = LIN_N, + prefix: str = "dyn_linear", + k: int = LIN_K, + group: int = LIN_GROUP, ) -> None: - from executorch.backends.webgpu.test.ops.quantized_linear.test_quantized_linear import ( + from executorch.backends.webgpu.test.ops.test_quantized_linear import ( _fp64_golden, _make_quantized_model, ) - model = _make_quantized_model(LIN_K, n, LIN_GROUP) - x = _ramp((LIN_MAXM, LIN_K)) + model = _make_quantized_model(k, n, group) + x = _ramp((LIN_MAXM, k)) m_dim = torch.export.Dim("m", min=1, max=LIN_MAXM) ep = torch.export.export(model, (x,), dynamic_shapes=({0: m_dim},)) - et = to_edge_transform_and_lower( - ep, partitioner=[VulkanPartitioner()] - ).to_executorch() - assert any( - d.id == "VulkanBackend" - for plan in et.executorch_program.execution_plan - for d in plan.delegates - ), "linear_q4gsw not delegated" + 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 [LIN_MAXM, 32, 1]: - xm = _ramp((m, LIN_K)) + xm = _ramp((m, k)) g = _fp64_golden(model, xm).astype(" None: + from executorch.backends.webgpu.test.ops.test_quantized_linear import ( + _fp64_golden, + _make_quantized_model, + ) + + model = _make_quantized_model(LIN_K, LIN_N, LIN_GROUP) + x = _ramp((m, LIN_K)) + ep = torch.export.export(model, (x,)) + et = _lower_fully_delegated(ep, prefix) + with open(os.path.join(out_dir, f"{prefix}.pte"), "wb") as f: + f.write(et.buffer) + x.detach().numpy().astype(" None: - from executorch.backends.webgpu.test.ops.sdpa.test_sdpa import ( + from executorch.backends.webgpu.test.ops.test_sdpa import ( _det_inputs, _golden, SdpaConfig, @@ -286,14 +337,7 @@ def cfg(s: int) -> "SdpaConfig": s_dim = torch.export.Dim("s", min=1, max=SD_MAXS) ds = ({1: s_dim}, {1: s_dim}, {1: s_dim}, None, None) ep = torch.export.export(model, (q, k, v, kc, vc), dynamic_shapes=ds) - et = to_edge_transform_and_lower( - ep, partitioner=[VulkanPartitioner()] - ).to_executorch() - assert any( - d.id == "VulkanBackend" - for plan in et.executorch_program.execution_plan - for d in plan.delegates - ), "sdpa not delegated" + et = _lower_fully_delegated(ep, "sdpa_dyn") with open(os.path.join(out_dir, "sdpa_dyn.pte"), "wb") as f: f.write(et.buffer) print("Exported sdpa_dyn.pte") @@ -315,6 +359,148 @@ def cfg(s: int) -> "SdpaConfig": print(f" golden sdpa_dyn S={s} (golden shape {tuple(g.shape)})") +def _export_combined_routes(out_dir: str) -> None: + from executorch.backends.webgpu.test.ops.test_quantized_linear import ( + _make_quantized_model, + ) + from executorch.backends.webgpu.test.ops.test_sdpa import ( + _det_inputs, + _golden, + SdpaConfig, + ) + from executorch.extension.llm.custom_ops import custom_ops # noqa: F401 + + class CombinedRoutes(torch.nn.Module): + def __init__(self): + super().__init__() + self.linear = _make_quantized_model(LIN_K, SD_HQ * SD_D, LIN_GROUP) + + def forward(self, x, q, k, v, k_cache, v_cache): + projected = self.linear(x).reshape(1, x.shape[0], SD_HQ, SD_D) + return torch.ops.llama.sdpa_with_kv_cache( + q + projected, + k, + v, + k_cache, + v_cache, + 0, + q.shape[1], + None, + 0.0, + True, + None, + ) + + model = CombinedRoutes().eval() + cfg = SdpaConfig("combined", SD_HQ, SD_HKV, SD_D, SD_MAXS, SD_CMAX, 0) + q, k, v, kc, vc = _det_inputs(cfg) + x = _ramp((SD_MAXS, LIN_K)) + s_dim = torch.export.Dim("s", min=1, max=SD_MAXS) + ep = torch.export.export( + model, + (x, q, k, v, kc, vc), + dynamic_shapes=( + {0: s_dim}, + {1: s_dim}, + {1: s_dim}, + {1: s_dim}, + None, + None, + ), + ) + et = _lower_fully_delegated(ep, "combined_routes") + with open(os.path.join(out_dir, "combined_routes.pte"), "wb") as f: + f.write(et.buffer) + print("Exported combined_routes.pte") + + for s in [SD_MAXS, 16, 1]: + live_cfg = SdpaConfig("combined", SD_HQ, SD_HKV, SD_D, s, SD_CMAX, 0) + q, k, v, kc, vc = _det_inputs(live_cfg) + x = _ramp((s, LIN_K)) + with torch.no_grad(): + projected = model.linear(x).reshape(1, s, SD_HQ, SD_D) + golden = _golden(live_cfg, q + projected, k, v, kc, vc) + base = os.path.join(out_dir, f"combined_routes.S{s}.") + for name, tensor in ( + ("x", x), + ("q", q), + ("k", k), + ("v", v), + ("kc", kc), + ("vc", vc), + ("golden", golden), + ): + tensor.detach().numpy().astype(" None: + from executorch.backends.webgpu.test.ops.test_sdpa import ( + _det_inputs, + _golden, + SdpaConfig, + SdpaModule, + ) + + hq, hkv, d, cmax, max_s = 8, 2, 132, 16, 16 + + def cfg(s): + return SdpaConfig("wide", hq, hkv, d, s, cmax, 0) + + model = SdpaModule(0) + q, k, v, kc, vc = _det_inputs(cfg(max_s)) + s_dim = torch.export.Dim("s", min=1, max=max_s) + ep = torch.export.export( + model, + (q, k, v, kc, vc), + dynamic_shapes=({1: s_dim}, {1: s_dim}, {1: s_dim}, None, None), + ) + et = _lower_fully_delegated(ep, "sdpa_wide") + with open(os.path.join(out_dir, "sdpa_wide.pte"), "wb") as f: + f.write(et.buffer) + for s in [max_s, 1]: + live = cfg(s) + q, k, v, kc, vc = _det_inputs(live) + golden = _golden(live, q, k, v, kc, vc) + for name, tensor in ( + ("q", q), + ("k", k), + ("v", v), + ("kc", kc), + ("vc", vc), + ("golden", golden), + ): + tensor.detach().numpy().astype(" None: + from executorch.backends.webgpu.test.ops.test_sdpa import ( + _det_inputs, + _golden, + SdpaConfig, + SdpaModule, + ) + + cfg = SdpaConfig(prefix, SD_HQ, SD_HKV, SD_D, s, SD_CMAX, 0) + inputs = _det_inputs(cfg) + ep = torch.export.export(SdpaModule(0), inputs) + et = _lower_fully_delegated(ep, prefix) + with open(os.path.join(out_dir, f"{prefix}.pte"), "wb") as f: + f.write(et.buffer) + golden = _golden(cfg, *inputs) + for name, tensor in zip(("q", "k", "v", "kc", "vc"), inputs): + tensor.detach().numpy().astype(" [N, EMBED] fp32. EMB_VOCAB = 64 EMB_DIM = 64 @@ -323,7 +509,7 @@ def cfg(s: int) -> "SdpaConfig": def _export_dynamic_embedding(out_dir: str) -> None: - from executorch.backends.webgpu.test.ops.embedding_q4gsw.test_embedding_q4gsw import ( + from executorch.backends.webgpu.test.ops.test_embedding_q4gsw import ( _make_quantized_model, _quant_params, Shape, @@ -334,14 +520,7 @@ def _export_dynamic_embedding(out_dir: str) -> None: idx_max = torch.arange(EMB_MAXN, dtype=torch.long) n_dim = torch.export.Dim("n", min=1, max=EMB_MAXN) ep = torch.export.export(qm, (idx_max,), dynamic_shapes=({0: n_dim},)) - et = to_edge_transform_and_lower( - ep, partitioner=[VulkanPartitioner()] - ).to_executorch() - assert any( - d.id == "VulkanBackend" - for plan in et.executorch_program.execution_plan - for d in plan.delegates - ), "embedding_q4gsw not delegated" + et = _lower_fully_delegated(ep, "emb_dyn") with open(os.path.join(out_dir, "emb_dyn.pte"), "wb") as f: f.write(et.buffer) print("Exported emb_dyn.pte") @@ -368,7 +547,7 @@ def _export_dynamic_embedding(out_dir: str) -> None: def _export_dynamic_rope(out_dir: str) -> None: - from executorch.backends.webgpu.test.ops.rope.test_rope import ( + from executorch.backends.webgpu.test.ops.test_rope import ( _golden, _inputs, Shape, @@ -381,14 +560,7 @@ def _export_dynamic_rope(out_dir: str) -> None: ep = torch.export.export( RotaryEmbedding().eval(), (xq, xk, fc, fs), dynamic_shapes=ds ) - et = to_edge_transform_and_lower( - ep, partitioner=[VulkanPartitioner()] - ).to_executorch() - assert any( - d.id == "VulkanBackend" - for plan in et.executorch_program.execution_plan - for d in plan.delegates - ), "apply_rotary_emb not delegated" + et = _lower_fully_delegated(ep, "rope_dyn") with open(os.path.join(out_dir, "rope_dyn.pte"), "wb") as f: f.write(et.buffer) print("Exported rope_dyn.pte") @@ -421,14 +593,7 @@ def _export_dynamic_select(out_dir: str) -> None: (_ramp((SEL_LEAD, 1, MAXS, HIDDEN)),), dynamic_shapes=({2: s_dim},), ) - et = to_edge_transform_and_lower( - ep, partitioner=[VulkanPartitioner()] - ).to_executorch() - assert any( - d.id == "VulkanBackend" - for plan in et.executorch_program.execution_plan - for d in plan.delegates - ), "select_copy not delegated" + et = _lower_fully_delegated(ep, "dyn_select") with open(os.path.join(out_dir, "dyn_select.pte"), "wb") as f: f.write(et.buffer) print("Exported dyn_select.pte") From 0619cdedcebcfaf9ece1c80952b9a1a46be6f5d4 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Tue, 21 Jul 2026 21:23:30 -0700 Subject: [PATCH 2/3] Update [ghstack-poisoned] --- .../test/ops/dynamic_shape/test_dynamic_shape_export.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) 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 57feaeeb414..635b877b8a4 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 @@ -547,11 +547,7 @@ def _export_dynamic_embedding(out_dir: str) -> None: def _export_dynamic_rope(out_dir: str) -> None: - from executorch.backends.webgpu.test.ops.test_rope import ( - _golden, - _inputs, - Shape, - ) + from executorch.backends.webgpu.test.ops.test_rope import _golden, _inputs, Shape from executorch.examples.models.llama.rope import RotaryEmbedding xq, xk, fc, fs = _inputs(Shape("dyn", 1, ROPE_MAXS, ROPE_NH, ROPE_NKV, ROPE_HD)) From 8bdee7cfc0533f79221bcaa3d27a4c6418c4ee71 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Wed, 29 Jul 2026 20:20:04 -0700 Subject: [PATCH 3/3] Update [ghstack-poisoned] --- backends/webgpu/runtime/WebGPUGraph.cpp | 91 +++++- backends/webgpu/runtime/WebGPUGraph.h | 86 ++++++ .../ops/quantized_linear/QuantizedLinear.cpp | 21 +- .../webgpu/runtime/ops/sigmoid/UnaryOp.cpp | 13 +- .../webgpu/runtime/ops/unary/Activations.cpp | 75 ++++- backends/webgpu/runtime/ops/unary/UnaryOp.cpp | 52 ++-- backends/webgpu/runtime/ops/unary/UnaryOp.h | 5 +- .../test/native/test_compute_dispatch.cpp | 285 ++++++++++++++++++ backends/webgpu/test/test_wgsl_codegen.py | 17 -- 9 files changed, 573 insertions(+), 72 deletions(-) diff --git a/backends/webgpu/runtime/WebGPUGraph.cpp b/backends/webgpu/runtime/WebGPUGraph.cpp index 5916434ab38..e275115e94b 100644 --- a/backends/webgpu/runtime/WebGPUGraph.cpp +++ b/backends/webgpu/runtime/WebGPUGraph.cpp @@ -426,6 +426,59 @@ size_t WebGPUGraph::add_compute_dispatch( return dispatch_index; } +size_t WebGPUGraph::add_dynamic_compute_dispatch_impl( + const WebGPUComputeDispatchDescriptor& descriptor, + int trigger_tensor_id, + std::function pick_grid) { + if (trigger_tensor_id < 0 || trigger_tensor_id >= num_values() || + get_value_type(trigger_tensor_id) != ValueType::Tensor) { + throw std::runtime_error( + "WebGPU dynamic dispatch: trigger must be a Tensor"); + } + if (!pick_grid) { + throw std::runtime_error("WebGPU dynamic dispatch: null grid picker"); + } + + const WebGPUDispatchGrid initial_grid = pick_grid(*this); + if (initial_grid.x == 0 || initial_grid.y == 0) { + throw std::runtime_error("WebGPU dynamic dispatch: zero grid"); + } + + WebGPUComputeDispatchDescriptor initial_descriptor = descriptor; + initial_descriptor.grid = initial_grid; + + // Reserve both vectors before creating GPU objects, then stage the sidecar. + // If GPU-object creation fails, removing the sidecar restores the graph; no + // operation that can fail remains after add_compute_dispatch succeeds. + const size_t new_size = dynamic_dispatch_grids_.size() + 1; + dynamic_dispatch_grids_.reserve(new_size); + pending_dynamic_dispatch_grids_.reserve(new_size); + + const size_t expected_index = dispatches_.size(); + dynamic_dispatch_grids_.push_back( + {expected_index, trigger_tensor_id, std::move(pick_grid)}); + try { + add_compute_dispatch(initial_descriptor); + } catch (...) { + dynamic_dispatch_grids_.pop_back(); + throw; + } + return expected_index; +} + +void WebGPUGraph::validate_dynamic_dispatch_route_ranges( + const std::vector& ranges) const { + for (const auto& dynamic_grid : dynamic_dispatch_grids_) { + for (const auto& range : ranges) { + if (range.begin <= dynamic_grid.dispatch_index && + dynamic_grid.dispatch_index < range.end) { + throw std::runtime_error( + "WebGPU dispatch cannot have both dynamic-grid and route ownership"); + } + } + } +} + void WebGPUGraph::update_symints_from_inputs( const std::vector& inputs) { for (const auto& src : symint_sources_) { @@ -562,11 +615,41 @@ void WebGPUGraph::propagate_resize() { pass++) { std::unordered_set processing; processing.swap(dirty_tensors_); - for (auto& hook : tensor_resize_hooks_) { - if (processing.count(hook.trigger_tensor_id) != 0) { - hook.fn(*this); + pending_dynamic_dispatch_grids_.clear(); + try { + for (auto& hook : tensor_resize_hooks_) { + if (processing.count(hook.trigger_tensor_id) != 0) { + hook.fn(*this); + } } - } + + // A hook or picker may fail, so compute and validate every affected grid + // before changing any dispatch. The graph-owned staging vector has + // capacity for every registered dynamic grid and is reused on execute. + for (const auto& dynamic_grid : dynamic_dispatch_grids_) { + if (processing.count(dynamic_grid.trigger_tensor_id) == 0) { + continue; + } + const WebGPUDispatchGrid grid = dynamic_grid.pick_grid(*this); + if (grid.x == 0 || grid.y == 0) { + throw std::runtime_error("WebGPU dynamic dispatch: zero grid"); + } + pending_dynamic_dispatch_grids_.push_back( + {dynamic_grid.dispatch_index, grid}); + } + } catch (...) { + pending_dynamic_dispatch_grids_.clear(); + // Keep both the current triggers and any cascaded outputs dirty so the + // caller can fix the hook or picker and retry without rebuilding. + dirty_tensors_.insert(processing.begin(), processing.end()); + throw; + } + for (const auto& pending : pending_dynamic_dispatch_grids_) { + auto& dispatch = dispatches_[pending.dispatch_index]; + dispatch.workgroup_count_x = pending.grid.x; + dispatch.workgroup_count_y = pending.grid.y; + } + pending_dynamic_dispatch_grids_.clear(); } if (!dirty_tensors_.empty()) { throw std::runtime_error( diff --git a/backends/webgpu/runtime/WebGPUGraph.h b/backends/webgpu/runtime/WebGPUGraph.h index 9d0bebcea36..45f21ccdc6a 100644 --- a/backends/webgpu/runtime/WebGPUGraph.h +++ b/backends/webgpu/runtime/WebGPUGraph.h @@ -12,10 +12,12 @@ #include #include +#include #include #include #include #include +#include #include #include @@ -250,9 +252,30 @@ class WebGPUGraph { // Per-SymInt resize hook; mirrors Vulkan DynamicDispatchNode::trigger_resize. void add_resize_hook(int symint_id, std::function fn) { + if (symint_id < 0 || symint_id >= num_values() || + get_value_type(symint_id) != ValueType::SymInt) { + throw std::runtime_error("WebGPU resize: trigger must be a SymInt"); + } + if (!fn) { + throw std::runtime_error("WebGPU resize: null SymInt resize hook"); + } resize_hooks_.push_back({symint_id, std::move(fn)}); } + template + void add_resize_hook( + int symint_id, + void (*fn)(WebGPUGraph&, const Context&), + Context context) { + if (fn == nullptr) { + throw std::runtime_error("WebGPU resize: null SymInt resize hook"); + } + add_resize_hook( + symint_id, [fn, context = std::move(context)](WebGPUGraph& graph) { + fn(graph, context); + }); + } + // Set a graph input's live dims (<= max) + dirty it; static path stays inert. void resize_input(int value_id, const std::vector& new_dims); // Set a tensor's live dims (an op resize hook calls this for its output to @@ -267,9 +290,31 @@ class WebGPUGraph { void add_tensor_resize_hook( int trigger_tensor_id, std::function fn) { + if (trigger_tensor_id < 0 || trigger_tensor_id >= num_values() || + get_value_type(trigger_tensor_id) != ValueType::Tensor) { + throw std::runtime_error("WebGPU resize: trigger must be a Tensor"); + } + if (!fn) { + throw std::runtime_error("WebGPU resize: null tensor resize hook"); + } tensor_resize_hooks_.push_back({trigger_tensor_id, std::move(fn)}); } + template + void add_tensor_resize_hook( + int trigger_tensor_id, + void (*fn)(WebGPUGraph&, const Context&), + Context context) { + if (fn == nullptr) { + throw std::runtime_error("WebGPU resize: null tensor resize hook"); + } + add_tensor_resize_hook( + trigger_tensor_id, + [fn, context = std::move(context)](WebGPUGraph& graph) { + fn(graph, context); + }); + } + // Run hooks for changed SymInts and tensors, then clear; call before execute. void propagate_resize(); @@ -283,6 +328,7 @@ class WebGPUGraph { size_t register_dispatch_route_group( const std::vector& ranges) { + validate_dynamic_dispatch_route_ranges(ranges); return dispatch_routes_.register_group( dispatches_.size(), ranges, [&](size_t i) { return dispatches_[i].kind == WebGPUDispatch::Kind::Compute; @@ -410,6 +456,23 @@ class WebGPUGraph { size_t add_compute_dispatch( const WebGPUComputeDispatchDescriptor& descriptor); + template + size_t add_dynamic_compute_dispatch( + const WebGPUComputeDispatchDescriptor& descriptor, + int trigger_tensor_id, + WebGPUDispatchGrid (*pick_grid)(const WebGPUGraph&, const Context&), + Context context) { + if (pick_grid == nullptr) { + throw std::runtime_error("WebGPU dynamic dispatch: null grid picker"); + } + return add_dynamic_compute_dispatch_impl( + descriptor, + trigger_tensor_id, + [pick_grid, context = std::move(context)](const WebGPUGraph& graph) { + return pick_grid(graph, context); + }); + } + WGPUShaderModule get_or_create_shader( const std::string& key, const char* wgsl_source); @@ -517,6 +580,29 @@ class WebGPUGraph { std::vector tensor_resize_hooks_; std::unordered_set dirty_tensors_; + // Dynamic grids are stored separately so ordinary dispatches remain compact. + // The graph owns each dispatch index and picker; ops only provide typed + // context and a named grid function. + struct DynamicDispatchGrid { + size_t dispatch_index; + int trigger_tensor_id; + std::function pick_grid; + }; + std::vector dynamic_dispatch_grids_; + + struct PendingDynamicDispatchGrid { + size_t dispatch_index; + WebGPUDispatchGrid grid; + }; + std::vector pending_dynamic_dispatch_grids_; + + size_t add_dynamic_compute_dispatch_impl( + const WebGPUComputeDispatchDescriptor& descriptor, + int trigger_tensor_id, + std::function pick_grid); + void validate_dynamic_dispatch_route_ranges( + const std::vector& ranges) const; + std::vector input_ids_; std::vector output_ids_; diff --git a/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp b/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp index 5f0cbb1ff20..16d78a6f923 100644 --- a/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp +++ b/backends/webgpu/runtime/ops/quantized_linear/QuantizedLinear.cpp @@ -21,7 +21,6 @@ #include #include -#include #include #include #include @@ -402,16 +401,7 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { use_steel, use_shmem_gemm); - WGPUBufferDescriptor uniform_desc = {}; - uniform_desc.size = sizeof(Q4gswParams); - 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(Q4gswParams)); - std::memcpy(mapped, &initial_state.params, sizeof(Q4gswParams)); - wgpuBufferUnmap(uniform_buffer); - graph.add_uniform_buffer_bytes(sizeof(Q4gswParams)); + WGPUBuffer params_buffer = graph.create_params_buffer(initial_state.params); const std::vector bindings = { {0, WGPUBufferBindingType_Storage, out.buffer, out.nbytes}, @@ -419,7 +409,7 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { {2, WGPUBufferBindingType_ReadOnlyStorage, weight.buffer, weight.nbytes}, {3, WGPUBufferBindingType_ReadOnlyStorage, scales.buffer, scales.nbytes}, {4, WGPUBufferBindingType_ReadOnlyStorage, bias_buffer, bias_size}, - {5, WGPUBufferBindingType_Uniform, uniform_buffer, sizeof(Q4gswParams)}, + {5, WGPUBufferBindingType_Uniform, params_buffer, sizeof(Q4gswParams)}, }; auto make_bundle = [&](const char* source, bool fixed_wg, uint32_t override_wg_size) { @@ -503,7 +493,7 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { use_shmem_gemm, dispatch_idx, route_group, - uniform_buffer](WebGPUGraph& g) { + params_buffer](WebGPUGraph& g) { const auto& d = g.cur_dims(in_id); const Q4gswExecutionState state = make_q4gsw_execution_state( g.device(), @@ -521,7 +511,7 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { use_steel, use_shmem_gemm); wgpuQueueWriteBuffer( - g.queue(), uniform_buffer, 0, &state.params, sizeof(state.params)); + g.queue(), params_buffer, 0, &state.params, sizeof(state.params)); if (use_dual_route) { g.select_dispatch_route( route_group, state.active_route, {state.active_grid}); @@ -532,9 +522,6 @@ void q4gsw_linear_impl(WebGPUGraph& graph, const std::vector& args) { } g.set_cur_dims(out_id, state.output_dims); }); - - // Graph owns it so the resize hook can rewrite it; freed in the dtor. - graph.own_uniform_buffer(uniform_buffer); } } // namespace diff --git a/backends/webgpu/runtime/ops/sigmoid/UnaryOp.cpp b/backends/webgpu/runtime/ops/sigmoid/UnaryOp.cpp index 49d2fbdb6bb..f7dca371ecf 100644 --- a/backends/webgpu/runtime/ops/sigmoid/UnaryOp.cpp +++ b/backends/webgpu/runtime/ops/sigmoid/UnaryOp.cpp @@ -7,6 +7,8 @@ */ #include +#include +#include #include #include @@ -17,12 +19,19 @@ namespace { void sigmoid_impl(WebGPUGraph& graph, const std::vector& args) { // aten.sigmoid.default args: [in, out] - add_unary_op(graph, args.at(0), args.at(1), "sigmoid", "sigmoid"); + add_unary_op( + graph, + args.at(0), + args.at(1), + kSigmoidWGSL, + kSigmoidWorkgroupSizeX, + "sigmoid"); } void relu_impl(WebGPUGraph& graph, const std::vector& args) { // aten.relu.default args: [in, out] - add_unary_op(graph, args.at(0), args.at(1), "relu", "relu"); + add_unary_op( + graph, args.at(0), args.at(1), kReluWGSL, kReluWorkgroupSizeX, "relu"); } } // namespace diff --git a/backends/webgpu/runtime/ops/unary/Activations.cpp b/backends/webgpu/runtime/ops/unary/Activations.cpp index eb71ba7bf8f..53b05e8b61b 100644 --- a/backends/webgpu/runtime/ops/unary/Activations.cpp +++ b/backends/webgpu/runtime/ops/unary/Activations.cpp @@ -8,6 +8,18 @@ #include #include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include #include #include @@ -34,64 +46,101 @@ float get_val_or_inf(WebGPUGraph& graph, int id, bool is_max) { } void abs_impl(WebGPUGraph& graph, const std::vector& args) { - add_unary_op(graph, args.at(0), args.at(1), "abs", "abs"); + add_unary_op( + graph, args.at(0), args.at(1), kAbsWGSL, kAbsWorkgroupSizeX, "abs"); } void exp_impl(WebGPUGraph& graph, const std::vector& args) { - add_unary_op(graph, args.at(0), args.at(1), "exp", "exp"); + add_unary_op( + graph, args.at(0), args.at(1), kExpWGSL, kExpWorkgroupSizeX, "exp"); } void sqrt_impl(WebGPUGraph& graph, const std::vector& args) { - add_unary_op(graph, args.at(0), args.at(1), "sqrt", "sqrt"); + add_unary_op( + graph, args.at(0), args.at(1), kSqrtWGSL, kSqrtWorkgroupSizeX, "sqrt"); } void rsqrt_impl(WebGPUGraph& graph, const std::vector& args) { - add_unary_op(graph, args.at(0), args.at(1), "rsqrt", "rsqrt"); + add_unary_op( + graph, args.at(0), args.at(1), kRsqrtWGSL, kRsqrtWorkgroupSizeX, "rsqrt"); } void sin_impl(WebGPUGraph& graph, const std::vector& args) { - add_unary_op(graph, args.at(0), args.at(1), "sin", "sin"); + add_unary_op( + graph, args.at(0), args.at(1), kSinWGSL, kSinWorkgroupSizeX, "sin"); } void cos_impl(WebGPUGraph& graph, const std::vector& args) { - add_unary_op(graph, args.at(0), args.at(1), "cos", "cos"); + add_unary_op( + graph, args.at(0), args.at(1), kCosWGSL, kCosWorkgroupSizeX, "cos"); } void tanh_impl(WebGPUGraph& graph, const std::vector& args) { - add_unary_op(graph, args.at(0), args.at(1), "tanh", "tanh"); + add_unary_op( + graph, args.at(0), args.at(1), kTanhWGSL, kTanhWorkgroupSizeX, "tanh"); } void round_impl(WebGPUGraph& graph, const std::vector& args) { - add_unary_op(graph, args.at(0), args.at(1), "round", "round"); + add_unary_op( + graph, args.at(0), args.at(1), kRoundWGSL, kRoundWorkgroupSizeX, "round"); } void neg_impl(WebGPUGraph& graph, const std::vector& args) { - add_unary_op(graph, args.at(0), args.at(1), "neg", "neg"); + add_unary_op( + graph, args.at(0), args.at(1), kNegWGSL, kNegWorkgroupSizeX, "neg"); } void hardswish_impl(WebGPUGraph& graph, const std::vector& args) { - add_unary_op(graph, args.at(0), args.at(1), "hardswish", "hardswish"); + add_unary_op( + graph, + args.at(0), + args.at(1), + kHardswishWGSL, + kHardswishWorkgroupSizeX, + "hardswish"); } void clamp_impl(WebGPUGraph& graph, const std::vector& args) { // aten.clamp.default args: [in, min, max, out]; min/max None -> +/-inf. const float lo = get_val_or_inf(graph, args.at(1), /*is_max=*/false); const float hi = get_val_or_inf(graph, args.at(2), /*is_max=*/true); - add_unary_op(graph, args.at(0), args.at(3), "clamp", "clamp", lo, hi); + add_unary_op( + graph, + args.at(0), + args.at(3), + kClampWGSL, + kClampWorkgroupSizeX, + "clamp", + lo, + hi); } void hardtanh_impl(WebGPUGraph& graph, const std::vector& args) { // aten.hardtanh.default args: [in, min_val, max_val, out]. const float lo = get_val_or_inf(graph, args.at(1), /*is_max=*/false); const float hi = get_val_or_inf(graph, args.at(2), /*is_max=*/true); - add_unary_op(graph, args.at(0), args.at(3), "clamp", "hardtanh", lo, hi); + add_unary_op( + graph, + args.at(0), + args.at(3), + kClampWGSL, + kClampWorkgroupSizeX, + "hardtanh", + lo, + hi); } void pow_scalar_impl(WebGPUGraph& graph, const std::vector& args) { // aten.pow.Tensor_Scalar args: [in, exponent, out]; exponent = min slot. const float exponent = get_val_or_inf(graph, args.at(1), /*is_max=*/false); add_unary_op( - graph, args.at(0), args.at(2), "pow_scalar", "pow_scalar", exponent); + graph, + args.at(0), + args.at(2), + kPowScalarWGSL, + kPowScalarWorkgroupSizeX, + "pow_scalar", + exponent); } } // namespace diff --git a/backends/webgpu/runtime/ops/unary/UnaryOp.cpp b/backends/webgpu/runtime/ops/unary/UnaryOp.cpp index 45ee6938b32..48564e74918 100644 --- a/backends/webgpu/runtime/ops/unary/UnaryOp.cpp +++ b/backends/webgpu/runtime/ops/unary/UnaryOp.cpp @@ -8,7 +8,6 @@ #include -#include #include #include @@ -35,7 +34,8 @@ void add_unary_op( WebGPUGraph& graph, int in_id, int out_id, - const char* shader_name, + const char* wgsl_source, + uint32_t wg_size_x, const char* op_name, float min, float max) { @@ -56,34 +56,47 @@ void add_unary_op( // Adaptive 1D->2D dispatch: wg=clamp(device,256) + 2D-spill past the 65535 // per-dim ceiling. The shader decodes idx via num_workgroups.x, so the live // count_x sets the stride at runtime (resize-safe, no override to re-bake). - const WebGPUShaderInfo& shader_info = get_webgpu_shader_info(shader_name); - uint32_t wg_size = - utils::clamp_workgroup_size(device, shader_info.workgroup_size_x); + uint32_t wg_size = utils::clamp_workgroup_size(device, wg_size_x); utils::WgCount workgroup_count = utils::compute_2d_workgroup_count(device, num_elements, wg_size, op_name); + WGPUConstantEntry wg_size_constant = utils::make_wg_size_constant(wg_size); + UnaryParams params = {}; params.num_elements = num_elements; params.min = min; params.max = max; - WGPUBuffer params_buffer = graph.create_params_buffer(params); + WGPUBuffer uniform_buffer = + utils::make_uniform(device, ¶ms, sizeof(UnaryParams)); + graph.add_uniform_buffer_bytes(sizeof(UnaryParams)); // input (read storage) + output (storage) + params. - WebGPUComputeDispatchDescriptor descriptor; - descriptor.shader_name = shader_name; - descriptor.kernel_name = op_name; - descriptor.bindings = { - {in_tensor.buffer, 0u, in_tensor.nbytes}, - {out_tensor.buffer, 0u, out_tensor.nbytes}, - {params_buffer, 0u, sizeof(UnaryParams)}, - }; - descriptor.constants = {{"wg_size", static_cast(wg_size)}}; - descriptor.grid = {workgroup_count.x, workgroup_count.y}; - const size_t dispatch_idx = graph.add_compute_dispatch(descriptor); + utils::ComputePipelineBundle bundle = utils::make_compute_pipeline( + device, + wgsl_source, + { + {0, + WGPUBufferBindingType_ReadOnlyStorage, + in_tensor.buffer, + in_tensor.nbytes}, + {1, + WGPUBufferBindingType_Storage, + out_tensor.buffer, + out_tensor.nbytes}, + {2, + WGPUBufferBindingType_Uniform, + uniform_buffer, + sizeof(UnaryParams)}, + }, + &wg_size_constant, + 1); + + const size_t dispatch_idx = graph.add_dispatch_2d( + bundle.pipeline, bundle.bind_group, workgroup_count.x, workgroup_count.y); // Dynamic shapes: rewrite full params (incl. min/max) + 2D dispatch count. - WGPUBuffer params_buf = params_buffer; + WGPUBuffer params_buf = uniform_buffer; graph.add_tensor_resize_hook( in_id, [in_id, out_id, wg_size, dispatch_idx, params_buf, min, max]( @@ -101,6 +114,9 @@ void add_unary_op( g.dispatch_at(dispatch_idx).workgroup_count_x = wgc.x; g.dispatch_at(dispatch_idx).workgroup_count_y = wgc.y; }); + + // Graph owns it so the resize hook can rewrite it; freed in the dtor. + graph.own_uniform_buffer(uniform_buffer); } } // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/unary/UnaryOp.h b/backends/webgpu/runtime/ops/unary/UnaryOp.h index 73b85ae0eba..897c6740b3c 100644 --- a/backends/webgpu/runtime/ops/unary/UnaryOp.h +++ b/backends/webgpu/runtime/ops/unary/UnaryOp.h @@ -10,6 +10,8 @@ #include +#include + namespace executorch::backends::webgpu { // Dummy min/max for no-param activations; mirrors Vulkan kDummyFloat. @@ -20,7 +22,8 @@ void add_unary_op( WebGPUGraph& graph, int in_id, int out_id, - const char* shader_name, + const char* wgsl_source, + uint32_t wg_size_x, const char* op_name, float min = kUnaryDummyFloat, float max = kUnaryDummyFloat); diff --git a/backends/webgpu/test/native/test_compute_dispatch.cpp b/backends/webgpu/test/native/test_compute_dispatch.cpp index 99c30cb7524..639a8ad8d41 100644 --- a/backends/webgpu/test/native/test_compute_dispatch.cpp +++ b/backends/webgpu/test/native/test_compute_dispatch.cpp @@ -368,6 +368,291 @@ TEST(WebGPUComputeDispatch, RejectsBindingRangeBeyondDawnBuffer) { wgpuBufferRelease(buffer); } +constexpr int kResizeQ = 0; +constexpr int kResizeK = 1; +constexpr int kResizeUnrelated = 2; +constexpr int kResizeCascade = 3; +constexpr int kResizeOutput = 4; +constexpr int kResizeSymInt = 5; + +void build_resize_test_graph(WebGPUGraph& graph) { + namespace vk = vkgraph; + ::flatbuffers::FlatBufferBuilder fbb; + std::vector<::flatbuffers::Offset> values; + const std::vector dims = {8, 8}; + for (int mem_obj_id = 0; mem_obj_id < 5; mem_obj_id++) { + values.push_back(vk::CreateVkValue( + fbb, + vk::GraphTypes::VkTensor, + vk::CreateVkTensorDirect( + fbb, + vk::VkDataType::FLOAT32, + &dims, + /*constant_id=*/-1, + mem_obj_id) + .Union())); + } + values.push_back(vk::CreateVkValue( + fbb, vk::GraphTypes::SymInt, vk::CreateSymInt(fbb, 0).Union())); + + const std::vector<::flatbuffers::Offset> chain; + const std::vector input_ids = { + kResizeQ, kResizeK, kResizeUnrelated}; + const std::vector output_ids = {kResizeOutput}; + const auto root = vk::CreateVkGraphDirect( + fbb, "0", &chain, &values, &input_ids, &output_ids); + vk::FinishVkGraphBuffer(fbb, root); + + graph.set_device(g_device); + graph.build(fbb.GetBufferPointer(), nullptr, 0, nullptr); +} + +WebGPUComputeDispatchDescriptor make_dynamic_test_descriptor( + WebGPUGraph& graph, + const char* kernel_name) { + const auto& input = graph.get_tensor(kResizeQ); + const auto& output = graph.get_tensor(kResizeOutput); + WGPUBuffer params = + graph.create_params_buffer(UnaryParams{64u, -1.0f, 1.0f, 0u}); + WebGPUComputeDispatchDescriptor descriptor; + descriptor.shader_name = "clamp"; + descriptor.kernel_name = kernel_name; + descriptor.bindings = { + {input.buffer, 0u, input.nbytes}, + {output.buffer, 0u, output.nbytes}, + {params, 0u, sizeof(UnaryParams)}}; + descriptor.constants = {{"wg_size", 64.0}}; + descriptor.grid = {99u, 97u}; + return descriptor; +} + +struct ResizeProbeContext { + int marker; + int* observed; + int* calls; +}; + +void record_resize_probe(WebGPUGraph&, const ResizeProbeContext& context) { + *context.observed = context.marker; + ++*context.calls; +} + +struct GridPickerContext { + int tensor_id; + uint32_t x_bias; + uint32_t y_bias; + int* calls; + const bool* fail; +}; + +WebGPUDispatchGrid pick_tensor_grid( + const WebGPUGraph& graph, + const GridPickerContext& context) { + ++*context.calls; + if (context.fail != nullptr && *context.fail) { + throw std::runtime_error("dynamic grid picker failure"); + } + const auto& dims = graph.cur_dims(context.tensor_id); + return { + static_cast(dims.at(0)) + context.x_bias, + static_cast(dims.at(1)) + context.y_bias}; +} + +struct CascadeContext { + int source_id; + int output_id; +}; + +void resize_cascade_output(WebGPUGraph& graph, const CascadeContext& context) { + graph.set_cur_dims(context.output_id, graph.cur_dims(context.source_id)); +} + +void expect_dispatch_grid( + WebGPUGraph& graph, + size_t dispatch_index, + uint32_t x, + uint32_t y) { + const auto& dispatch = graph.dispatch_at(dispatch_index); + EXPECT_EQ(dispatch.workgroup_count_x, x); + EXPECT_EQ(dispatch.workgroup_count_y, y); +} + +TEST(WebGPUResizeHooks, TypedRegistrationOwnsContextCopies) { + WebGPUGraph graph; + build_resize_test_graph(graph); + int tensor_observed = 0; + int tensor_calls = 0; + int symint_observed = 0; + int symint_calls = 0; + { + ResizeProbeContext tensor_context = {17, &tensor_observed, &tensor_calls}; + ResizeProbeContext symint_context = {23, &symint_observed, &symint_calls}; + graph.add_tensor_resize_hook(kResizeQ, record_resize_probe, tensor_context); + graph.add_resize_hook(kResizeSymInt, record_resize_probe, symint_context); + tensor_context.marker = 101; + symint_context.marker = 103; + } + + graph.resize_input(kResizeQ, {7, 8}); + graph.propagate_resize(); + EXPECT_EQ(tensor_observed, 17); + EXPECT_EQ(tensor_calls, 1); + EXPECT_EQ(symint_observed, 0); + EXPECT_EQ(symint_calls, 0); + + graph.set_symint(kResizeSymInt, 9); + graph.propagate_resize(); + EXPECT_EQ(tensor_calls, 1); + EXPECT_EQ(symint_observed, 23); + EXPECT_EQ(symint_calls, 1); +} + +TEST(WebGPUDynamicDispatch, InitializesAndIsolatesTriggeredGrids) { + WebGPUGraph graph; + build_resize_test_graph(graph); + int q_calls = 0; + int k_calls = 0; + GridPickerContext q_context = {kResizeQ, 1u, 2u, &q_calls, nullptr}; + const GridPickerContext k_context = {kResizeK, 3u, 4u, &k_calls, nullptr}; + const size_t q_dispatch = graph.add_dynamic_compute_dispatch( + make_dynamic_test_descriptor(graph, "dynamic_q"), + kResizeQ, + pick_tensor_grid, + q_context); + const size_t k_dispatch = graph.add_dynamic_compute_dispatch( + make_dynamic_test_descriptor(graph, "dynamic_k"), + kResizeK, + pick_tensor_grid, + k_context); + q_context.x_bias = 101u; + q_context.y_bias = 103u; + expect_dispatch_grid(graph, q_dispatch, 9u, 10u); + expect_dispatch_grid(graph, k_dispatch, 11u, 12u); + EXPECT_EQ(q_calls, 1); + EXPECT_EQ(k_calls, 1); + q_calls = 0; + k_calls = 0; + + graph.resize_input(kResizeUnrelated, {7, 7}); + graph.propagate_resize(); + EXPECT_EQ(q_calls, 0); + EXPECT_EQ(k_calls, 0); + + graph.resize_input(kResizeQ, {4, 3}); + graph.propagate_resize(); + expect_dispatch_grid(graph, q_dispatch, 5u, 5u); + expect_dispatch_grid(graph, k_dispatch, 11u, 12u); + EXPECT_EQ(q_calls, 1); + EXPECT_EQ(k_calls, 0); + + graph.resize_input(kResizeQ, {2, 5}); + graph.propagate_resize(); + expect_dispatch_grid(graph, q_dispatch, 3u, 7u); + EXPECT_EQ(q_calls, 2); + graph.resize_input(kResizeQ, {2, 5}); + graph.propagate_resize(); + EXPECT_EQ(q_calls, 2); + + graph.resize_input(kResizeK, {6, 1}); + graph.propagate_resize(); + expect_dispatch_grid(graph, k_dispatch, 9u, 5u); + EXPECT_EQ(k_calls, 1); +} + +TEST(WebGPUDynamicDispatch, HandlesCascadesAndStagesPickerFailures) { + { + WebGPUGraph graph; + build_resize_test_graph(graph); + int same_pass_calls = 0; + int cascade_pass_calls = 0; + graph.add_tensor_resize_hook( + kResizeQ, + resize_cascade_output, + CascadeContext{kResizeQ, kResizeCascade}); + const size_t same_pass_dispatch = graph.add_dynamic_compute_dispatch( + make_dynamic_test_descriptor(graph, "dynamic_same_pass"), + kResizeQ, + pick_tensor_grid, + GridPickerContext{kResizeCascade, 5u, 7u, &same_pass_calls, nullptr}); + const size_t cascade_pass_dispatch = graph.add_dynamic_compute_dispatch( + make_dynamic_test_descriptor(graph, "dynamic_cascade_pass"), + kResizeCascade, + pick_tensor_grid, + GridPickerContext{ + kResizeCascade, 9u, 11u, &cascade_pass_calls, nullptr}); + same_pass_calls = 0; + cascade_pass_calls = 0; + graph.resize_input(kResizeQ, {4, 6}); + graph.propagate_resize(); + EXPECT_EQ(same_pass_calls, 1); + EXPECT_EQ(cascade_pass_calls, 1); + expect_dispatch_grid(graph, same_pass_dispatch, 9u, 13u); + expect_dispatch_grid(graph, cascade_pass_dispatch, 13u, 17u); + } + + WebGPUGraph graph; + build_resize_test_graph(graph); + int first_calls = 0; + int second_calls = 0; + bool second_fails = false; + const size_t first = graph.add_dynamic_compute_dispatch( + make_dynamic_test_descriptor(graph, "dynamic_first"), + kResizeQ, + pick_tensor_grid, + GridPickerContext{kResizeQ, 1u, 2u, &first_calls, nullptr}); + const size_t second = graph.add_dynamic_compute_dispatch( + make_dynamic_test_descriptor(graph, "dynamic_second"), + kResizeQ, + pick_tensor_grid, + GridPickerContext{kResizeQ, 3u, 4u, &second_calls, &second_fails}); + expect_dispatch_grid(graph, first, 9u, 10u); + expect_dispatch_grid(graph, second, 11u, 12u); + second_fails = true; + graph.resize_input(kResizeQ, {4, 3}); + EXPECT_THROW(graph.propagate_resize(), std::runtime_error); + expect_dispatch_grid(graph, first, 9u, 10u); + expect_dispatch_grid(graph, second, 11u, 12u); +} + +TEST(WebGPUDynamicDispatch, RejectsRouteOverlapWithoutPoisoningRegistry) { + WebGPUGraph graph; + build_resize_test_graph(graph); + int calls = 0; + bool initial_fails = true; + const size_t dispatches_before_failure = graph.num_dispatches(); + EXPECT_THROW( + graph.add_dynamic_compute_dispatch( + make_dynamic_test_descriptor(graph, "dynamic_initial_failure"), + kResizeQ, + pick_tensor_grid, + GridPickerContext{kResizeQ, 0u, 0u, &calls, &initial_fails}), + std::runtime_error); + EXPECT_EQ(graph.num_dispatches(), dispatches_before_failure); + + const size_t dynamic = graph.add_dynamic_compute_dispatch( + make_dynamic_test_descriptor(graph, "dynamic_route_guard"), + kResizeQ, + pick_tensor_grid, + GridPickerContext{kResizeQ, 0u, 0u, &calls, nullptr}); + ASSERT_EQ(dynamic, 0u); + graph.add_dispatch(WebGPUDispatch{}); + graph.add_dispatch(WebGPUDispatch{}); + + EXPECT_THROW( + graph.register_dispatch_route_group({{0, 1}, {1, 2}}), + std::runtime_error); + calls = 0; + graph.resize_input(kResizeQ, {7, 6}); + graph.propagate_resize(); + EXPECT_EQ(calls, 1); + expect_dispatch_grid(graph, dynamic, 7u, 6u); + const size_t group = graph.register_dispatch_route_group({{1, 2}, {2, 3}}); + EXPECT_EQ(group, 0u); + graph.select_dispatch_route(group, 1, {{13u, 17u}}); + expect_dispatch_grid(graph, 1u, 0u, 0u); + expect_dispatch_grid(graph, 2u, 13u, 17u); +} + TEST(WebGPUExecution, FullySuppressedPlanPerformsNoQueueSubmission) { WebGPUGraph graph; const WebGPUExecutionPlan plan; diff --git a/backends/webgpu/test/test_wgsl_codegen.py b/backends/webgpu/test/test_wgsl_codegen.py index 3816589608c..fe8340d6528 100644 --- a/backends/webgpu/test/test_wgsl_codegen.py +++ b/backends/webgpu/test/test_wgsl_codegen.py @@ -97,23 +97,6 @@ def test_registry_rejects_duplicate_names(self) -> None: with self.assertRaisesRegex(ValueError, "duplicate shader registry name"): g.render_registry([entry, entry]) - def test_unary_builder_uses_graph_dispatch_descriptor(self) -> None: - unary = (g.BACKEND_ROOT / "runtime/ops/unary/UnaryOp.cpp").read_text() - self.assertIn("get_webgpu_shader_info(shader_name)", unary) - self.assertIn("workgroup_size_x", unary) - self.assertIn("graph.create_params_buffer", unary) - self.assertIn("graph.add_compute_dispatch", unary) - self.assertIn("workgroup_count.x, workgroup_count.y", unary) - self.assertIn("workgroup_count_x = wgc.x", unary) - self.assertIn("workgroup_count_y = wgc.y", unary) - for legacy in ( - "utils::make_uniform", - "utils::make_compute_pipeline", - "graph.add_uniform_buffer_bytes", - "graph.own_uniform_buffer", - ): - self.assertNotIn(legacy, unary) - def test_symbol_base(self) -> None: self.assertEqual(g.symbol_base("binary_add"), "BinaryAdd") self.assertEqual(