Skip to content

Commit e5b236c

Browse files
committed
[ExecuTorch][WebGPU] SDPA test suite: replay + dynamic input_pos + in-graph KV cache
Pull Request resolved: pytorch#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: 393014582 @exported-using-ghexport Differential Revision: [D107595144](https://our.internmc.facebook.com/intern/diff/D107595144/)
1 parent 7bb5a3d commit e5b236c

4 files changed

Lines changed: 1663 additions & 0 deletions

File tree

backends/webgpu/scripts/test_webgpu_native_ci.sh

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -75,6 +75,21 @@ 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+
# Non-fatal: a failed sdpa export makes the required 4k/8k configs hard-fail in
79+
# webgpu_native_test below (precise per-config error), so don't exit/mask here.
80+
$PYTHON_EXECUTABLE -c "
81+
from executorch.backends.webgpu.test.ops.sdpa.test_sdpa import (
82+
export_all_sdpa_models,
83+
export_replay_sequences,
84+
export_dynamic_decode,
85+
export_incache_decode,
86+
)
87+
export_all_sdpa_models('/tmp')
88+
export_replay_sequences('/tmp')
89+
export_dynamic_decode('/tmp')
90+
export_incache_decode('/tmp')
91+
" || echo "WARN: sdpa export failed; required 4k/8k configs will FAIL in webgpu_native_test"
92+
7893
# ── Configure (Dawn-only: no -DWEBGPU_IMPL; Dawn is the sole backend) ─────────
7994
echo "=== Configure WebGPU native tests on Dawn ==="
8095
rm -rf "${BUILD_DIR}"

0 commit comments

Comments
 (0)