diff --git a/backends/webgpu/runtime/WebGPUDispatchMath.h b/backends/webgpu/runtime/WebGPUDispatchMath.h index 9b87990aa8f..561bebc6b87 100644 --- a/backends/webgpu/runtime/WebGPUDispatchMath.h +++ b/backends/webgpu/runtime/WebGPUDispatchMath.h @@ -96,8 +96,9 @@ constexpr bool should_record_q4gsw_dual_route( constexpr bool should_record_sdpa_dual_route( bool fd_eligible, - bool has_dynamic_sequence) { - return fd_eligible && has_dynamic_sequence; + bool has_dynamic_sequence, + bool has_dynamic_position) { + return fd_eligible && (has_dynamic_sequence || has_dynamic_position); } constexpr bool is_q4gsw_bk64_eligible( diff --git a/backends/webgpu/runtime/WebGPUGraph.cpp b/backends/webgpu/runtime/WebGPUGraph.cpp index e729b3926a0..859e01e7a40 100644 --- a/backends/webgpu/runtime/WebGPUGraph.cpp +++ b/backends/webgpu/runtime/WebGPUGraph.cpp @@ -1340,19 +1340,86 @@ void WebGPUGraph::copy_inputs(const std::vector& inputs) { } } +#ifdef WGPU_BACKEND_ENABLE_PROFILING +// Profiling/attestation only; never compiled into a production build. Written +// during WebGPUGraph::execute without synchronization: the attestation +// harnesses that read them run one graph on one thread, so no atomics or +// locking are needed. To support concurrent profiled execution, make these +// per-instance behind a whole-record mutex (per-field atomics would not cover +// the conflict check's read-modify-write across both globals). +uint32_t g_last_route_mask = 0; +uint32_t g_last_route_conflict_count = 0; +#endif // WGPU_BACKEND_ENABLE_PROFILING + namespace { +#ifdef WGPU_BACKEND_ENABLE_PROFILING +constexpr uint32_t kRoutePrefill = 1u << 0; +constexpr uint32_t kRouteK16 = 1u << 1; +constexpr uint32_t kRouteMaterializedAttention = 1u << 2; +constexpr uint32_t kRouteT0Steel = 1u << 3; +constexpr uint32_t kRouteT1Bk64 = 1u << 4; +constexpr uint32_t kRouteT1Bk64Qkv = 1u << 5; +constexpr uint32_t kRouteT2PairedGateUp = 1u << 6; +constexpr uint32_t kRouteFusedSwiGlu = 1u << 7; +constexpr uint32_t kRouteGenericFallback = 1u << 8; +// bit 1u << 9 is intentionally reserved (a retired route) and left unused. +constexpr uint32_t kRouteFlashDecoding = 1u << 10; +constexpr uint32_t kRouteK16CausalBound = 1u << 11; +constexpr uint32_t kRouteBicolSubgroup = 1u << 12; +#endif // WGPU_BACKEND_ENABLE_PROFILING + // Bench gate: compiled out unless WGPU_BACKEND_ENABLE_PROFILING; then the // WEBGPU_TIMESTAMP_QUERY env var enables per-pass GPU timestamp queries. bool should_timestamp_query() { #ifdef WGPU_BACKEND_ENABLE_PROFILING - static const bool enabled = std::getenv("WEBGPU_TIMESTAMP_QUERY") != nullptr; - return enabled; + return std::getenv("WEBGPU_TIMESTAMP_QUERY") != nullptr; #else return false; #endif } } // namespace +#ifdef WGPU_BACKEND_ENABLE_PROFILING +void WebGPUGraph::record_active_route(const std::string& kernel_name) { + uint32_t bits = 0; + if (kernel_name == "sdpa_streaming_attention_k16_causal_bound") { + bits = kRoutePrefill | kRouteK16CausalBound; + } else if (kernel_name == "sdpa_streaming_attention_k16") { + bits = kRoutePrefill | kRouteK16; + } else if ( + kernel_name.rfind("sdpa_compute_", 0) == 0 || + kernel_name == "sdpa_softmax") { + bits = kRoutePrefill | kRouteMaterializedAttention; + } else if (kernel_name == "fd_split" || kernel_name == "fd_reduce") { + bits = kRouteFlashDecoding; + } else if (kernel_name == "linear_q4gsw_coop4_bicol_subgroup") { + bits = kRouteBicolSubgroup; + } else if (kernel_name.rfind("linear_q4gsw_bk64_qkv", 0) == 0) { + bits = kRouteT1Bk64Qkv; + } else if (kernel_name.rfind("linear_q4gsw_bk64", 0) == 0) { + bits = kRouteT1Bk64; + } else if (kernel_name.rfind("linear_q4gsw_paired_gate_up", 0) == 0) { + bits = kRouteT2PairedGateUp; + } else if (kernel_name == "silu_mul_fused") { + bits = kRouteFusedSwiGlu; + } else if (kernel_name.rfind("linear_q4gsw_steel", 0) == 0) { + bits = kRouteT0Steel; + } else if (kernel_name.rfind("linear_q4gsw", 0) == 0) { + bits = kRouteGenericFallback; + } + + constexpr uint32_t kAttentionRoutes = kRouteK16 | kRouteK16CausalBound | + kRouteMaterializedAttention | kRouteFlashDecoding; + const uint32_t new_attention = bits & kAttentionRoutes; + const uint32_t prior_attention = g_last_route_mask & kAttentionRoutes; + if (new_attention != 0 && prior_attention != 0 && + (new_attention & prior_attention) == 0) { + ++g_last_route_conflict_count; + } + g_last_route_mask |= bits; +} +#endif // WGPU_BACKEND_ENABLE_PROFILING + WebGPUExecutionPlan WebGPUGraph::make_execution_plan( const WebGPUGraphExecutionOptions& options) const { const size_t n = dispatches_.size(); @@ -1378,6 +1445,10 @@ WebGPUExecutionPlan WebGPUGraph::make_execution_plan( } size_t WebGPUGraph::execute(const WebGPUExecutionPlan& plan) { +#ifdef WGPU_BACKEND_ENABLE_PROFILING + g_last_route_mask = 0; + g_last_route_conflict_count = 0; +#endif // WGPU_BACKEND_ENABLE_PROFILING const size_t n = dispatches_.size(); const size_t chunk = execute_config_.chunk_size; if (plan.copy_outputs.size() != output_copies_.size()) { @@ -1443,6 +1514,9 @@ size_t WebGPUGraph::execute(const WebGPUExecutionPlan& plan) { dispatch.copy_nbytes); continue; } +#ifdef WGPU_BACKEND_ENABLE_PROFILING + record_active_route(dispatch.kernel_name); +#endif // WGPU_BACKEND_ENABLE_PROFILING WGPUComputePassDescriptor pass_desc = {}; #ifdef WGPU_BACKEND_ENABLE_PROFILING // tw must outlive BeginComputePass (the descriptor points at it). @@ -1530,6 +1604,9 @@ size_t WebGPUGraph::execute(const WebGPUExecutionPlan& plan) { dispatches_[i].copy_nbytes); continue; } +#ifdef WGPU_BACKEND_ENABLE_PROFILING + record_active_route(dispatches_[i].kernel_name); +#endif // WGPU_BACKEND_ENABLE_PROFILING WGPUComputePassDescriptor pass_desc = {}; WGPUComputePassEncoder pass = wgpuCommandEncoderBeginComputePass(encoder, &pass_desc); diff --git a/backends/webgpu/runtime/WebGPUGraph.h b/backends/webgpu/runtime/WebGPUGraph.h index a70f099dc8e..292f2ee0cf4 100644 --- a/backends/webgpu/runtime/WebGPUGraph.h +++ b/backends/webgpu/runtime/WebGPUGraph.h @@ -546,6 +546,10 @@ class WebGPUGraph { } private: +#ifdef WGPU_BACKEND_ENABLE_PROFILING + void record_active_route(const std::string& kernel_name); +#endif // WGPU_BACKEND_ENABLE_PROFILING + bool kv_f16_ = false; std::unordered_set kv_cache_ids_; WebGPUGraphConfig config_; @@ -668,4 +672,9 @@ class WebGPUGraph { size_t uniform_buffer_bytes_ = 0; }; +#ifdef WGPU_BACKEND_ENABLE_PROFILING +extern uint32_t g_last_route_mask; +extern uint32_t g_last_route_conflict_count; +#endif // WGPU_BACKEND_ENABLE_PROFILING + } // namespace executorch::backends::webgpu diff --git a/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp b/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp index f36a84cc54f..9a51f853caf 100644 --- a/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp +++ b/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp @@ -596,8 +596,8 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { const bool dynamic_sequence = graph.tensor_has_dynamic_dims(q_id) || graph.tensor_has_dynamic_dims(k_id) || graph.tensor_has_dynamic_dims(v_id); - const bool dual_route = - utils::should_record_sdpa_dual_route(fd_eligible, dynamic_sequence); + const bool dual_route = utils::should_record_sdpa_dual_route( + fd_eligible, dynamic_sequence, dynamic_pos); const bool record_k16 = k16_eligible && (dual_route || !initial_state.use_fd); const bool record_materialized = !k16_eligible && (dual_route || !initial_state.use_fd); diff --git a/backends/webgpu/test/native/test_dispatch_2d.cpp b/backends/webgpu/test/native/test_dispatch_2d.cpp index 67ce74659eb..c9ce9648b43 100644 --- a/backends/webgpu/test/native/test_dispatch_2d.cpp +++ b/backends/webgpu/test/native/test_dispatch_2d.cpp @@ -217,9 +217,10 @@ TEST(DispatchRoute, RecordsEligibleQ4AndDynamicSdpaAlternates) { EXPECT_TRUE(should_record_q4gsw_dual_route(32, true, true, false)); EXPECT_TRUE(should_record_q4gsw_dual_route(32, true, false, true)); - EXPECT_FALSE(should_record_sdpa_dual_route(true, false)); - EXPECT_FALSE(should_record_sdpa_dual_route(false, true)); - EXPECT_TRUE(should_record_sdpa_dual_route(true, true)); + EXPECT_FALSE(should_record_sdpa_dual_route(true, false, false)); + EXPECT_FALSE(should_record_sdpa_dual_route(false, true, true)); + EXPECT_TRUE(should_record_sdpa_dual_route(true, true, false)); + EXPECT_TRUE(should_record_sdpa_dual_route(true, false, true)); } TEST(WebGPUGraphConfig, ParsesExactBooleanCompileOption) {