Skip to content

Commit f4052f9

Browse files
[ExecuTorch][WebGPU] SDPA test suite: replay + dynamic input_pos + in-graph KV cache
Pull Request resolved: #20087 Adds the WebGPU SDPA test coverage as its own diff, stacked on the SDPA op (which already carries the dynamic-`input_pos` consumption) and the SymInt mechanism below it: multi-step prefill->mt->decode replay, runtime-dynamic `input_pos` (autoregressive decode), and an in-graph mutable KV cache, each compared against a torch `F.scaled_dot_product_attention` golden. - `test/ops/sdpa/test_sdpa.py`: `ReplaySeq`/`REPLAY_SEQS` + per-step replay export/golden; `DynamicSdpaModule` + `export_dynamic_decode` (one `.pte`, `input_pos` supplied at runtime as a SymInt); `DecodeCacheModule` + `export_incache_decode` (KV cache as `register_buffer` mutable buffers, so the cache persists in-graph and forward() feeds only the new token + `input_pos`). - `test/test_webgpu_native.cpp`: `test_sdpa_replay`, `test_sdpa_dynamic_decode` (+ negative control: a pinned `input_pos` diverges), `test_sdpa_incache_decode` (+ static control: a fresh Module per step diverges, proving in-graph accumulation is real), `test_symint_roundtrip`, `test_resize_hook`; shared per-element tolerance `sdpa_within_tol` (abs 1e-4 OR rel 1e-3). - `test/test_build_webgpu.sh`: export the replay / dynamic / in-graph-cache models for the native test. Authored with assistance from Claude. ghstack-source-id: 391626190 @exported-using-ghexport Differential Revision: [D107595144](https://our.internmc.facebook.com/intern/diff/D107595144/)
1 parent 2e2d743 commit f4052f9

4 files changed

Lines changed: 1620 additions & 0 deletions

File tree

backends/webgpu/scripts/test_webgpu_native_ci.sh

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -75,6 +75,19 @@ export_update_cache_replay('${UPDATE_CACHE_DIR}')
7575
export_update_cache_negative('${UPDATE_CACHE_DIR}')
7676
" || { echo "WARN: update_cache export failed; skipping update_cache native test"; UPDATE_CACHE_OK=0; }
7777

78+
$PYTHON_EXECUTABLE -c "
79+
from executorch.backends.webgpu.test.ops.sdpa.test_sdpa import (
80+
export_all_sdpa_models,
81+
export_replay_sequences,
82+
export_dynamic_decode,
83+
export_incache_decode,
84+
)
85+
export_all_sdpa_models('/tmp')
86+
export_replay_sequences('/tmp')
87+
export_dynamic_decode('/tmp')
88+
export_incache_decode('/tmp')
89+
" || echo "WARN: sdpa export failed; webgpu_native_test sdpa cases will report missing .pte"
90+
7891
# ── Configure (Dawn-only: no -DWEBGPU_IMPL; Dawn is the sole backend) ─────────
7992
echo "=== Configure WebGPU native tests on Dawn ==="
8093
rm -rf "${BUILD_DIR}"

0 commit comments

Comments
 (0)