[ExecuTorch][WebGPU] GPU timestamp query profiling for SDPA

Pull Request resolved: #20167

Add a faithful re-port of Vulkan's `vkapi::QueryPool` (`backends/vulkan/runtime/vk_api/QueryPool.{h,cpp}`) so a bench can read true on-GPU per-kernel time, isolated from submit/readback latency — the basis for comparing the WGSL SDPA kernels against the Vulkan reference. Opt-in via the `WEBGPU_TIMESTAMP_QUERY` env var; off by default, so the production `execute()` path is byte-identical.

`WebGPUQueryPool` mirrors the Vulkan `ShaderDuration` data model and the ticks->ns conversion exactly. Three deviations are forced by the WebGPU API (not unforced divergences): per-dispatch bracketing uses a compute-pass `timestampWrites` descriptor (begin/end-of-pass) since WebGPU has no mid-encoder `writeTimestamp`; results are read via `resolveQuerySet` + buffer map (no host-side `vkGetQueryPoolResults`); and the `TimestampQuery` capability is requested as an explicit device feature (fail-open if the adapter lacks it). `WebGPUGraph::execute()` brackets each compute pass when the pool is active; chained `update_cache`/QK/softmax/AV dispatches carry a `kernel_name` label for attribution.

Co-authored-with Claude.
ghstack-source-id: 392065610
@exported-using-ghexport

Differential Revision: [D107678235](https://our.internmc.facebook.com/intern/diff/D107678235/)

Author: JCNTH

backends/webgpu/runtime/ops/sdpa/Sdpa.cpp

@@ -156,7 +156,8 @@ void build_dispatch(
     uint64_t uniform_size,
     uint32_t workgroup_count_x,
     uint32_t wg_size,
-    bool retain_uniform = false) {
+    bool retain_uniform = false,
+    const char* kernel_name = "") {
   WGPUDevice device = graph.device();
 
   WGPUShaderSourceWGSL wgsl_desc = {};
@@ -227,7 +228,7 @@ void build_dispatch(
   bg_desc.entries = bg_entries;
   WGPUBindGroup bind_group = wgpuDeviceCreateBindGroup(device, &bg_desc);
 
-  graph.add_dispatch({pipeline, bind_group, workgroup_count_x});
+  graph.add_dispatch({pipeline, bind_group, workgroup_count_x, kernel_name});
 
   wgpuShaderModuleRelease(shader);
   wgpuBindGroupLayoutRelease(bgl);
@@ -269,7 +270,8 @@ static WGPUBuffer record_update_cache_dispatch(
       sizeof(uc),
       wgc,
       uc_wg,
-      dynamic_pos);
+      dynamic_pos,
+      "update_cache");
   return ubuf;
 }
 
@@ -473,7 +475,8 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector<int>& args) {
         sizeof(p),
         wgc,
         qk_wg,
-        dynamic_pos);
+        dynamic_pos,
+        "sdpa_compute_attn_weights");
     qk_buf = ubuf;
     qk_idx = graph.num_dispatches() - 1;
   }
@@ -496,7 +499,8 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector<int>& args) {
         sizeof(p),
         wgc,
         0,
-        dynamic_pos);
+        dynamic_pos,
+        "sdpa_softmax");
     softmax_buf = ubuf;
   }
 
@@ -521,7 +525,8 @@ void sdpa_with_kv_cache_impl(WebGPUGraph& graph, const std::vector<int>& args) {
         sizeof(p),
         wgc,
         av_wg,
-        dynamic_pos);
+        dynamic_pos,
+        "sdpa_compute_out");
     av_buf = ubuf;
   }
 

backends/webgpu/test/test_webgpu_native.cpp

@@ -1122,6 +1122,129 @@ static bool test_resize_hook(const std::string& blob_path) {
   return true;
 }
 
+// Capacity-overrun must throw; runs without a device or TimestampQuery.
+static bool test_query_pool_overrun_throws() {
+  printf("\n--- Test: WebGPUQueryPool capacity-overrun guard ---\n");
+  WebGPUQueryPool qp;
+  try {
+    qp.reset(1);
+  } catch (const std::exception&) {
+    printf("PASS: reset beyond capacity throws\n");
+    return true;
+  }
+  printf("FAIL: reset beyond capacity did not throw\n");
+  return false;
+}
+
+// WebGPUQueryPool roundtrip: time a probe pass; assert non-zero GPU duration.
+static bool test_query_pool_roundtrip(const WebGPUContext& ctx) {
+  printf("\n--- Test: WebGPUQueryPool roundtrip ---\n");
+  if (!ctx.timestamp_supported) {
+    printf("SKIP: adapter lacks TimestampQuery feature\n");
+    return true;
+  }
+  WGPUDevice device = ctx.device;
+
+  // Probe loop iterates enough to burn a measurable, non-zero GPU duration.
+  const char* kProbeWGSL =
+      "@group(0) @binding(0) var<storage, read_write> out: array<f32>;\n"
+      "@compute @workgroup_size(64)\n"
+      "fn main(@builtin(global_invocation_id) gid: vec3<u32>) {\n"
+      "  var acc = 0.0;\n"
+      "  for (var i = 0u; i < 8192u; i = i + 1u) {\n"
+      "    acc = acc + f32(i) * 1.000001;\n"
+      "  }\n"
+      "  out[gid.x] = acc;\n"
+      "}\n";
+
+  WGPUShaderSourceWGSL wgsl_desc = {};
+  wgsl_desc.chain.sType = WGPUSType_ShaderSourceWGSL;
+  wgsl_desc.code = {kProbeWGSL, WGPU_STRLEN};
+  WGPUShaderModuleDescriptor shader_desc = {};
+  shader_desc.nextInChain = &wgsl_desc.chain;
+  WGPUShaderModule shader = wgpuDeviceCreateShaderModule(device, &shader_desc);
+
+  WGPUBindGroupLayoutEntry bgl_entry = {};
+  bgl_entry.binding = 0;
+  bgl_entry.visibility = WGPUShaderStage_Compute;
+  bgl_entry.buffer.type = WGPUBufferBindingType_Storage;
+  WGPUBindGroupLayoutDescriptor bgl_desc = {};
+  bgl_desc.entryCount = 1;
+  bgl_desc.entries = &bgl_entry;
+  WGPUBindGroupLayout bgl = wgpuDeviceCreateBindGroupLayout(device, &bgl_desc);
+
+  WGPUPipelineLayoutDescriptor pl_desc = {};
+  pl_desc.bindGroupLayoutCount = 1;
+  pl_desc.bindGroupLayouts = &bgl;
+  WGPUPipelineLayout pl = wgpuDeviceCreatePipelineLayout(device, &pl_desc);
+
+  WGPUComputePipelineDescriptor pipe_desc = {};
+  pipe_desc.layout = pl;
+  pipe_desc.compute.module = shader;
+  pipe_desc.compute.entryPoint = {"main", WGPU_STRLEN};
+  WGPUComputePipeline pipe =
+      wgpuDeviceCreateComputePipeline(device, &pipe_desc);
+
+  WGPUBufferDescriptor obd = {};
+  obd.size = 64 * sizeof(float);
+  obd.usage = WGPUBufferUsage_Storage;
+  WGPUBuffer out_buf = wgpuDeviceCreateBuffer(device, &obd);
+
+  WGPUBindGroupEntry bg_entry = {};
+  bg_entry.binding = 0;
+  bg_entry.buffer = out_buf;
+  bg_entry.size = obd.size;
+  WGPUBindGroupDescriptor bg_desc = {};
+  bg_desc.layout = bgl;
+  bg_desc.entryCount = 1;
+  bg_desc.entries = &bg_entry;
+  WGPUBindGroup bg = wgpuDeviceCreateBindGroup(device, &bg_desc);
+
+  WebGPUQueryPool qp;
+  qp.initialize(device, 1);
+  qp.reset(1);
+
+  WGPUCommandEncoder enc = wgpuDeviceCreateCommandEncoder(device, nullptr);
+  WGPUPassTimestampWrites tw = qp.writes_for(0);
+  WGPUComputePassDescriptor pass_desc = {};
+  pass_desc.timestampWrites = &tw;
+  WGPUComputePassEncoder pass =
+      wgpuCommandEncoderBeginComputePass(enc, &pass_desc);
+  wgpuComputePassEncoderSetPipeline(pass, pipe);
+  wgpuComputePassEncoderSetBindGroup(pass, 0, bg, 0, nullptr);
+  wgpuComputePassEncoderDispatchWorkgroups(pass, 1, 1, 1);
+  wgpuComputePassEncoderEnd(pass);
+  wgpuComputePassEncoderRelease(pass);
+  qp.record(0, "probe", {1, 1, 1}, {64, 1, 1});
+  qp.resolve(enc);
+  WGPUCommandBuffer cmd = wgpuCommandEncoderFinish(enc, nullptr);
+  wgpuQueueSubmit(ctx.queue, 1, &cmd);
+  wgpuCommandBufferRelease(cmd);
+  wgpuCommandEncoderRelease(enc);
+
+  qp.extract_results(ctx.instance);
+
+  wgpuBufferRelease(out_buf);
+  wgpuComputePipelineRelease(pipe);
+  wgpuPipelineLayoutRelease(pl);
+  wgpuBindGroupLayoutRelease(bgl);
+  wgpuBindGroupRelease(bg);
+  wgpuShaderModuleRelease(shader);
+
+  if (qp.results().size() != 1) {
+    printf("FAIL: expected 1 duration, got %zu\n", qp.results().size());
+    return false;
+  }
+  const uint64_t dur = qp.results()[0].execution_duration_ns;
+  printf("  probe duration: %llu ns\n", (unsigned long long)dur);
+  if (dur == 0) {
+    printf("FAIL: probe duration is zero (expected monotonic non-zero)\n");
+    return false;
+  }
+  printf("PASS: WebGPUQueryPool roundtrip -- non-zero GPU kernel duration\n");
+  return true;
+}
+
 int main(int argc, char** argv) {
   std::string model_path = "webgpu_add_test.pte";
   if (argc > 1) {
@@ -1163,7 +1286,9 @@ int main(int argc, char** argv) {
   set_default_webgpu_context(&ctx);
   printf("WebGPU device acquired (native)\n");
 
-  bool ok = test_single_add(model_path);
+  bool ok = test_query_pool_overrun_throws();
+  ok = test_query_pool_roundtrip(ctx) && ok;
+  ok = test_single_add(model_path) && ok;
 
   if (!chained_model_path.empty()) {
     ok = test_chained_add(chained_model_path) && ok;