From 651340f1588a4a12468ca74288bb1b6c3741fdd5 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Tue, 21 Jul 2026 21:11:23 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- backends/webgpu/runtime/WebGPUGraph.cpp | 74 ++++++++++++++++++- backends/webgpu/runtime/WebGPUGraph.h | 9 +++ backends/webgpu/runtime/WebGPUUtils.h | 5 +- backends/webgpu/runtime/ops/sdpa/Sdpa.cpp | 3 +- .../webgpu/test/native/test_dispatch_2d.cpp | 7 +- 5 files changed, 90 insertions(+), 8 deletions(-) diff --git a/backends/webgpu/runtime/WebGPUGraph.cpp b/backends/webgpu/runtime/WebGPUGraph.cpp index 2da6198a95c..6ac261bfdd5 100644 --- a/backends/webgpu/runtime/WebGPUGraph.cpp +++ b/backends/webgpu/runtime/WebGPUGraph.cpp @@ -1726,20 +1726,84 @@ void WebGPUGraph::copy_inputs(const std::vector& inputs) { } } +#ifdef WGPU_BACKEND_ENABLE_PROFILING +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; +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 + void WebGPUGraph::execute(const WebGPUGraphExecutionOptions& options) { +#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; std::vector enabled_dispatches(n, true); @@ -1808,6 +1872,9 @@ void WebGPUGraph::execute(const WebGPUGraphExecutionOptions& options) { 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). @@ -1893,6 +1960,9 @@ void WebGPUGraph::execute(const WebGPUGraphExecutionOptions& options) { 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 8fae1eae056..40e459750ca 100644 --- a/backends/webgpu/runtime/WebGPUGraph.h +++ b/backends/webgpu/runtime/WebGPUGraph.h @@ -384,6 +384,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_; bool f16_accumulate_gemm_ = false; @@ -491,4 +495,9 @@ class WebGPUGraph { void add_swiglu_fused_dispatch(int gate_id, int up_id, int out_id); }; +#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/WebGPUUtils.h b/backends/webgpu/runtime/WebGPUUtils.h index a623e923762..d3d17188c9d 100644 --- a/backends/webgpu/runtime/WebGPUUtils.h +++ b/backends/webgpu/runtime/WebGPUUtils.h @@ -78,8 +78,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/ops/sdpa/Sdpa.cpp b/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp index a516fda9046..b5468854ee3 100644 --- a/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp +++ b/backends/webgpu/runtime/ops/sdpa/Sdpa.cpp @@ -686,7 +686,8 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector& args) { 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); + 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 = diff --git a/backends/webgpu/test/native/test_dispatch_2d.cpp b/backends/webgpu/test/native/test_dispatch_2d.cpp index b80fdb3b377..8211633e2b7 100644 --- a/backends/webgpu/test/native/test_dispatch_2d.cpp +++ b/backends/webgpu/test/native/test_dispatch_2d.cpp @@ -212,9 +212,10 @@ TEST(DispatchRoute, RecordsAlternatesOnlyForDynamicEligibleGraphs) { 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)); + 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(DispatchRoute, Bk64RequiresExactLlamaShapeAndCapabilities) {