[ExecuTorch][WebGPU] GPU timestamp query profiling for SDPA#20167
[ExecuTorch][WebGPU] GPU timestamp query profiling for SDPA#20167JulianCloudNTH wants to merge 6 commits into
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/20167
Note: Links to docs will display an error until the docs builds have been completed. ❌ 1 New Failure, 2 Unrelated FailuresAs of commit ebf063e with merge base 5526971 ( NEW FAILURE - The following job has failed:
FLAKY - The following job failed but was likely due to flakiness present on trunk:
BROKEN TRUNK - The following job failed but was present on the merge base:👉 Rebase onto the `viable/strict` branch to avoid these failures
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
This PR needs a
|
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: 391669549 @exported-using-ghexport Differential Revision: [D107678235](https://our.internmc.facebook.com/intern/diff/D107678235/)
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: 391669549 @exported-using-ghexport Differential Revision: [D107678235](https://our.internmc.facebook.com/intern/diff/D107678235/)
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: 391741952 @exported-using-ghexport Differential Revision: [D107678235](https://our.internmc.facebook.com/intern/diff/D107678235/)
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: 391801048 @exported-using-ghexport Differential Revision: [D107678235](https://our.internmc.facebook.com/intern/diff/D107678235/)
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/)
Pull Request resolved: #20167 SDPA-specific instrumentation layered on the general GPU-timestamp infrastructure (companion diff below): tag each fused SDPA dispatch with its `kernel_name` so the `WebGPUQueryPool` can attribute on-GPU time to the attention stage that produced it. `sdpa_with_kv_cache` runs four chained dispatches — `update_cache` -> QK (`attn_weights`) -> softmax -> AV (`compute_out`); `WebGPUGraph::execute()` brackets each compute pass with a timestamp when the pool is active, and this diff labels each dispatch so the per-pass durations map back to the right stage. Opt-in via the `WEBGPU_TIMESTAMP_QUERY` env var; off by default, so the production `execute()` path is byte-identical. This is the per-kernel hook a forthcoming SDPA kernel benchmark will read; the benchmark itself (and any comparative numbers) is a separate follow-up. Co-authored with Claude. ghstack-source-id: 392093463 @exported-using-ghexport Differential Revision: [D107678235](https://our.internmc.facebook.com/intern/diff/D107678235/)
Pull Request resolved: #20167 SDPA-specific instrumentation layered on the general GPU-timestamp infrastructure (companion diff below): tag each fused SDPA dispatch with its `kernel_name` so the `WebGPUQueryPool` can attribute on-GPU time to the attention stage that produced it. `sdpa_with_kv_cache` runs four chained dispatches — `update_cache` -> QK (`attn_weights`) -> softmax -> AV (`compute_out`); `WebGPUGraph::execute()` brackets each compute pass with a timestamp when the pool is active, and this diff labels each dispatch so the per-pass durations map back to the right stage. Opt-in via the `WEBGPU_TIMESTAMP_QUERY` env var; off by default, so the production `execute()` path is byte-identical. This is the per-kernel hook a forthcoming SDPA kernel benchmark will read; the benchmark itself (and any comparative numbers) is a separate follow-up. Co-authored with Claude. ghstack-source-id: 392093463 @exported-using-ghexport Differential Revision: [D107678235](https://our.internmc.facebook.com/intern/diff/D107678235/)
Pull Request resolved: #20167 SDPA-specific instrumentation layered on the general GPU-timestamp infrastructure (companion diff below): tag each fused SDPA dispatch with its `kernel_name` so the `WebGPUQueryPool` can attribute on-GPU time to the attention stage that produced it. `sdpa_with_kv_cache` runs four chained dispatches — `update_cache` -> QK (`attn_weights`) -> softmax -> AV (`compute_out`); `WebGPUGraph::execute()` brackets each compute pass with a timestamp when the pool is active, and this diff labels each dispatch so the per-pass durations map back to the right stage. Opt-in via the `WEBGPU_TIMESTAMP_QUERY` env var; off by default, so the production `execute()` path is byte-identical. This is the per-kernel hook a forthcoming SDPA kernel benchmark will read; the benchmark itself (and any comparative numbers) is a separate follow-up. Co-authored with Claude. ghstack-source-id: 392093463 @exported-using-ghexport Differential Revision: [D107678235](https://our.internmc.facebook.com/intern/diff/D107678235/)
Stack from ghstack (oldest at bottom):
SDPA-specific instrumentation layered on the general GPU-timestamp infrastructure (companion diff below): tag each fused SDPA dispatch with its
kernel_nameso theWebGPUQueryPoolcan attribute on-GPU time to the attention stage that produced it.sdpa_with_kv_cacheruns four chained dispatches —update_cache-> QK (attn_weights) -> softmax -> AV (compute_out);WebGPUGraph::execute()brackets each compute pass with a timestamp when the pool is active, and this diff labels each dispatch so the per-pass durations map back to the right stage. Opt-in via theWEBGPU_TIMESTAMP_QUERYenv var; off by default, so the productionexecute()path is byte-identical. This is the per-kernel hook a forthcoming SDPA kernel benchmark will read; the benchmark itself (and any comparative numbers) is a separate follow-up.Co-authored with Claude.
@exported-using-ghexport
Differential Revision: D107678235
Differential Revision: D107678235