From ed97e6f5d27594a511c6d637988ee8dfec9ab465 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Fri, 7 Aug 2026 11:56:52 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- backends/webgpu/CMakeLists.txt | 17 + .../webgpu/scripts/test_webgpu_native_ci.sh | 42 +- backends/webgpu/test/BUCK | 54 ++ .../webgpu/test/native/RequiredDevicePolicy.h | 17 + .../webgpu/test/native/test_dispatch_2d.cpp | 38 ++ .../webgpu/test/native/test_dynamic_shape.cpp | 225 ++++++- .../webgpu/test/native/test_slice_chain.cpp | 122 ++++ .../webgpu/test/native/test_update_cache.cpp | 320 ++++++++- .../test/native/test_update_cache_state.cpp | 125 ++++ backends/webgpu/test/op_tests/cases.py | 22 +- .../webgpu/test/op_tests/op_test_driver.cpp | 5 +- .../test_dynamic_shape_export.py | 122 +++- backends/webgpu/test/ops/test_et_vk_sdpa.py | 30 +- .../webgpu/test/ops/test_rope_hf_single.py | 44 ++ backends/webgpu/test/ops/test_update_cache.py | 316 ++++++++- backends/webgpu/test/targets.bzl | 17 + .../webgpu/test/test_native_ci_contract.py | 40 ++ examples/models/gemma4/BUCK | 42 ++ .../gemma4/generate_target_prefill_oracle.py | 403 ++++++++++++ .../models/gemma4/target_prefill_contract.py | 350 ++++++++++ examples/models/gemma4/tests/targets.bzl | 64 ++ .../gemma4/tests/test_export_partitioners.py | 53 ++ .../models/gemma4/tests/test_export_smoke.py | 39 ++ .../tests/test_gemma4_plain_wasm_contract.py | 183 ++++++ .../tests/test_gemma4_sdpa_host_contract.py | 191 ++++++ .../test_generate_target_prefill_oracle.py | 618 ++++++++++++++++++ .../tests/test_selected_row_cross_decoder.py | 78 +++ .../tests/test_webgpu_artifact_manifest.py | 405 ++++++++++++ 28 files changed, 3944 insertions(+), 38 deletions(-) create mode 100644 backends/webgpu/test/native/RequiredDevicePolicy.h create mode 100644 backends/webgpu/test/native/test_slice_chain.cpp create mode 100644 backends/webgpu/test/native/test_update_cache_state.cpp create mode 100644 backends/webgpu/test/ops/test_rope_hf_single.py create mode 100644 examples/models/gemma4/generate_target_prefill_oracle.py create mode 100644 examples/models/gemma4/target_prefill_contract.py create mode 100644 examples/models/gemma4/tests/test_export_partitioners.py create mode 100644 examples/models/gemma4/tests/test_export_smoke.py create mode 100644 examples/models/gemma4/tests/test_gemma4_plain_wasm_contract.py create mode 100644 examples/models/gemma4/tests/test_gemma4_sdpa_host_contract.py create mode 100644 examples/models/gemma4/tests/test_generate_target_prefill_oracle.py create mode 100644 examples/models/gemma4/tests/test_selected_row_cross_decoder.py create mode 100644 examples/models/gemma4/tests/test_webgpu_artifact_manifest.py diff --git a/backends/webgpu/CMakeLists.txt b/backends/webgpu/CMakeLists.txt index 0298fdd5a3c..a112f30d275 100644 --- a/backends/webgpu/CMakeLists.txt +++ b/backends/webgpu/CMakeLists.txt @@ -261,6 +261,10 @@ if(EXECUTORCH_BUILD_WEBGPU_TEST) webgpu_update_cache_test test/native/test_update_cache.cpp ) target_link_libraries(webgpu_update_cache_test PRIVATE GTest::gtest) + target_include_directories( + webgpu_update_cache_test + PRIVATE "${EXECUTORCH_ROOT}/third-party/json/single_include" + ) add_webgpu_native_test( webgpu_dynamic_shape_test test/native/test_dynamic_shape.cpp ) @@ -275,6 +279,19 @@ if(EXECUTORCH_BUILD_WEBGPU_TEST) target_link_libraries( webgpu_dispatch_2d_test PRIVATE GTest::gtest GTest::gtest_main ) + add_webgpu_native_test( + webgpu_update_cache_state_test + test/native/test_update_cache_state.cpp + ) + target_link_libraries( + webgpu_update_cache_state_test PRIVATE GTest::gtest GTest::gtest_main + ) + add_webgpu_native_test( + webgpu_slice_chain_test test/native/test_slice_chain.cpp + ) + target_link_libraries( + webgpu_slice_chain_test PRIVATE GTest::gtest GTest::gtest_main + ) add_executable( webgpu_execution_options_test test/native/test_execution_options.cpp runtime/WebGPUExecutionOptions.cpp diff --git a/backends/webgpu/scripts/test_webgpu_native_ci.sh b/backends/webgpu/scripts/test_webgpu_native_ci.sh index 0d5d7a9e799..fcf81066cdc 100644 --- a/backends/webgpu/scripts/test_webgpu_native_ci.sh +++ b/backends/webgpu/scripts/test_webgpu_native_ci.sh @@ -61,6 +61,35 @@ run_with_required_device() { fi } +run_required_gtests() { + local output + if ! output="$(run_with_required_device "$@" 2>&1)"; then + printf '%s\n' "${output}" + return 1 + fi + printf '%s\n' "${output}" + + local test_name + for test_name in \ + DynamicShape.SliceCrosses2dDispatchBoundary \ + DynamicShape.CatCrosses2dDispatchBoundary \ + DynamicShape.SliceDualStoreWritesBothDestinations; do + if ! grep -Eq "^\\[ OK \\] ${test_name}( \\([0-9]+ ms\\))?$" \ + <<<"${output}"; then + echo "ERROR: required WebGPU test did not pass: ${test_name}" >&2 + return 1 + fi + done + if ! grep -Fxq '[ PASSED ] 3 tests.' <<<"${output}"; then + echo "ERROR: required WebGPU run did not pass exactly three tests" >&2 + return 1 + fi + if grep -Eq '^\\[ SKIPPED \\]' <<<"${output}"; then + echo "ERROR: required WebGPU run skipped a test" >&2 + return 1 + fi +} + DISPATCH_ORDER_DIR="/tmp/dispatch_order" UPDATE_CACHE_DIR="/tmp/update_cache" INDEX_DIR="/tmp/index" @@ -138,10 +167,11 @@ from executorch.backends.webgpu.test.ops.index.test_index import export_all_inde export_all_index_models('${INDEX_DIR}') " -$PYTHON_EXECUTABLE -c " +WEBGPU_TEST_HEAVY=1 $PYTHON_EXECUTABLE -c " from executorch.backends.webgpu.test.ops.dynamic_shape.test_dynamic_shape_export import export_dynamic_shape_cases export_dynamic_shape_cases('${DYNAMIC_SHAPE_DIR}') " +require_file "${DYNAMIC_SHAPE_DIR}/dyn_cat_2d.pte" $PYTHON_EXECUTABLE -c " from executorch.backends.webgpu.test.ops.test_sdpa import ( @@ -221,6 +251,9 @@ run_with_required_device env WEBGPU_TEST_SDPA_DIR=/tmp/ \ "${BIN_DIR}/webgpu_dispatch_order_test" "${DISPATCH_ORDER_DIR}" "${BIN_DIR}/webgpu_index_test" "${INDEX_DIR}" "${BIN_DIR}/webgpu_dynamic_shape_test" "${DYNAMIC_SHAPE_DIR}" +run_required_gtests env WEBGPU_REQUIRE_DEVICE=1 WEBGPU_TEST_HEAVY=1 \ + "${BIN_DIR}/webgpu_dynamic_shape_test" "${DYNAMIC_SHAPE_DIR}" \ + --gtest_filter=DynamicShape.SliceCrosses2dDispatchBoundary:DynamicShape.CatCrosses2dDispatchBoundary:DynamicShape.SliceDualStoreWritesBothDestinations "${BIN_DIR}/webgpu_scratch_buffer_test" "${BIN_DIR}/webgpu_dispatch_2d_test" "${BIN_DIR}/webgpu_compute_dispatch_test" @@ -238,4 +271,11 @@ $PYTHON_EXECUTABLE -m executorch.backends.webgpu.test.op_tests.generate_op_tests --output "${OP_TEST_DIR}" cmake --build "${BUILD_DIR}" --target webgpu_op_test -j"${NPROC}" "${BIN_DIR}/webgpu_op_test" --manifest "${OP_TEST_DIR}/manifest.json" +CAT_2D_TEST_DIR="/tmp/webgpu_cat_2d_test" +WEBGPU_TEST_HEAVY=1 $PYTHON_EXECUTABLE \ + -m executorch.backends.webgpu.test.op_tests.generate_op_tests \ + --output "${CAT_2D_TEST_DIR}" --ops cat +run_with_required_device env WEBGPU_REQUIRE_DEVICE=1 \ + "${BIN_DIR}/webgpu_op_test" \ + --manifest "${CAT_2D_TEST_DIR}/manifest.json" echo "=== WebGPU op-test framework on Dawn: passed ===" diff --git a/backends/webgpu/test/BUCK b/backends/webgpu/test/BUCK index bcddaa0f66b..079fb77cb81 100644 --- a/backends/webgpu/test/BUCK +++ b/backends/webgpu/test/BUCK @@ -19,6 +19,21 @@ fbcode_target( ], ) +fbcode_target( + _kind = python_unittest, + name = "test_et_vk_sdpa", + srcs = [ + "ops/test_et_vk_sdpa.py", + ], + deps = [ + "//caffe2:torch", + "//executorch/backends/vulkan:custom_ops_lib", + "//executorch/backends/vulkan/partitioner:vulkan_partitioner", + "//executorch/backends/vulkan:vulkan_preprocess", + "//executorch/exir:lib", + ], +) + non_fbcode_target( _kind = runtime.python_test, name = "test_webgpu_artifact_manifest", @@ -42,6 +57,18 @@ fbcode_target( ], ) +fbcode_target( + _kind = python_unittest, + name = "test_rope_hf_single", + srcs = [ + "ops/test_rope_hf_single.py", + ], + deps = [ + "//caffe2:torch", + "//executorch/backends/vulkan:vulkan_preprocess", + ], +) + fbcode_target( _kind = runtime.python_library, name = "tester", @@ -60,3 +87,30 @@ fbcode_target( "//executorch/backends/webgpu/scripts:webgpu_artifact_manifest", ], ) + +fbcode_target( + _kind = python_unittest, + name = "test_wgsl_codegen", + srcs = ["test_wgsl_codegen.py"], + deps = [ + "fbsource//third-party/pypi/pyyaml:pyyaml", + ], +) + +fbcode_target( + _kind = runtime.python_test, + name = "test_update_cache", + srcs = ["ops/test_update_cache.py"], + preload_deps = [ + "//executorch/extension/llm/custom_ops:custom_ops_aot_lib", + "//executorch/extension/llm/custom_ops:custom_ops_aot_py", + ], + deps = [ + "//caffe2:torch", + "//executorch/backends/vulkan/partitioner:vulkan_partitioner", + "//executorch/backends/vulkan/serialization:lib", + "//executorch/backends/vulkan:vulkan_preprocess", + "//executorch/examples/models/gemma4:webgpu_support", + "//executorch/exir:lib", + ], +) diff --git a/backends/webgpu/test/native/RequiredDevicePolicy.h b/backends/webgpu/test/native/RequiredDevicePolicy.h new file mode 100644 index 00000000000..565500f6a42 --- /dev/null +++ b/backends/webgpu/test/native/RequiredDevicePolicy.h @@ -0,0 +1,17 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +namespace executorch::backends::webgpu { + +inline int required_device_failure_exit_code(bool required) { + return required ? 1 : 0; +} + +} // namespace executorch::backends::webgpu diff --git a/backends/webgpu/test/native/test_dispatch_2d.cpp b/backends/webgpu/test/native/test_dispatch_2d.cpp index 7c5a5e141f3..a37858a991a 100644 --- a/backends/webgpu/test/native/test_dispatch_2d.cpp +++ b/backends/webgpu/test/native/test_dispatch_2d.cpp @@ -12,6 +12,9 @@ #include #include #include +#include +#include +#include #include @@ -20,6 +23,9 @@ #include #include +using executorch::backends::webgpu::required_device_failure_exit_code; +using executorch::backends::webgpu::set_cat_dispatch_grid; +using executorch::backends::webgpu::set_slice_dispatch_grid; using executorch::backends::webgpu::WebGPUDispatch; using executorch::backends::webgpu::WebGPUGraph; using executorch::backends::webgpu::utils::DispatchRange; @@ -65,6 +71,38 @@ TEST(DispatchFold, ThrowsWhenNeeds3rdDimension) { EXPECT_ANY_THROW(fold_workgroup_count_2d(kMax * kMax + 1u, kMax, "test")); } +TEST(SliceDispatchGrid, RestoresBothDimensionsAcrossResize) { + WebGPUGraph graph; + const size_t dispatch_index = graph.add_dispatch(WebGPUDispatch{}); + + for (const WgCount grid : + {WgCount{256u, 256u}, WgCount{65535u, 1u}, WgCount{256u, 256u}}) { + set_slice_dispatch_grid(graph, dispatch_index, grid); + EXPECT_EQ(graph.dispatch_at(dispatch_index).workgroup_count_x, grid.x); + EXPECT_EQ(graph.dispatch_at(dispatch_index).workgroup_count_y, grid.y); + } +} + +TEST(CatDispatchGrid, RestoresBothDimensionsAcrossResize) { + WebGPUGraph graph; + const size_t dispatch_index = graph.add_dispatch(WebGPUDispatch{}); + + for (const WgCount grid : + {WgCount{257u, 256u}, WgCount{65535u, 1u}, WgCount{257u, 256u}}) { + set_cat_dispatch_grid(graph, dispatch_index, grid); + EXPECT_EQ(graph.dispatch_at(dispatch_index).workgroup_count_x, grid.x); + EXPECT_EQ(graph.dispatch_at(dispatch_index).workgroup_count_y, grid.y); + } +} + +TEST(RequiredDevicePolicy, DefaultDeviceFailureRemainsASkip) { + EXPECT_EQ(required_device_failure_exit_code(false), 0); +} + +TEST(RequiredDevicePolicy, RequiredDeviceFailureIsAnError) { + EXPECT_NE(required_device_failure_exit_code(true), 0); +} + void expect_grid(const WgCount& grid, uint32_t x, uint32_t y) { EXPECT_EQ(grid.x, x); EXPECT_EQ(grid.y, y); diff --git a/backends/webgpu/test/native/test_dynamic_shape.cpp b/backends/webgpu/test/native/test_dynamic_shape.cpp index cf16cd7e87c..24f9ea84c42 100644 --- a/backends/webgpu/test/native/test_dynamic_shape.cpp +++ b/backends/webgpu/test/native/test_dynamic_shape.cpp @@ -27,6 +27,8 @@ #include #include +#include +#include #include #include #include @@ -169,6 +171,19 @@ void check_conv1d(Module& module, int length) { constexpr int kGeluOld1dDispatchCap = 4 * 64 * 65535; constexpr int kGelu2dDispatchBoundary = kGeluOld1dDispatchCap + 1; constexpr int kGeluPatternSize = 257; +constexpr int kSliceInputColumns = 65; +constexpr int kSliceOutputColumns = 64; +constexpr int kSlice2dRows = 65536; +constexpr int kSlice1dRows = 65535; +constexpr int kSlicePatternSize = 4093; +constexpr int kSliceDualRows = 512; +constexpr int kSliceDualInputColumns = 8961; +constexpr int kSliceDualOutputColumns = 8960; +constexpr int kCatInputColumns = 65; +constexpr int kCatOutputColumns = 66; +constexpr int kCat2dRows = 64527; +constexpr int kCat1dRows = 64526; +constexpr int kCatPatternSize = 4091; void check_gelu_2d(Module& module, int elements) { std::array golden = {}; @@ -198,6 +213,163 @@ void check_gelu_2d(Module& module, int elements) { << " max_err=" << error; } +float slice_value(size_t linear_index, int salt) { + return static_cast( + (linear_index * 17 + static_cast(salt) * 101) % + kSlicePatternSize); +} + +#ifdef WGPU_BACKEND_ENABLE_PROFILING +bool slice_profile_available() { + const auto* context = get_default_webgpu_context(); + return std::getenv("WEBGPU_TIMESTAMP_QUERY") != nullptr && + context != nullptr && context->timestamp_supported && + context->querypool != nullptr; +} + +void expect_slice_profile(int rows) { + const auto* context = get_default_webgpu_context(); + ASSERT_NE(context, nullptr); + ASSERT_NE(context->querypool, nullptr); + const auto& profile = context->querypool->results(); + const auto slice = + std::find_if(profile.begin(), profile.end(), [](const auto& duration) { + return duration.kernel_name == "slice"; + }); + ASSERT_NE(slice, profile.end()) << "missing slice dispatch for rows=" << rows; + const std::array expected = rows == kSlice2dRows + ? std::array{256, 256, 1} + : std::array{65535, 1, 1}; + EXPECT_EQ(slice->global_wg, expected) << "slice rows=" << rows; +} +#endif + +void check_slice_2d(Module& module, int rows, int salt) { + const size_t input_numel = static_cast(rows) * kSliceInputColumns; + std::vector input(input_numel); + for (size_t i = 0; i < input.size(); ++i) { + input[i] = slice_value(i, salt); + } + auto tensor = make_tensor_ptr({rows, kSliceInputColumns}, std::move(input)); + auto result = module.forward({EValue(tensor)}); + ASSERT_TRUE( + result.ok() && result.get().size() == 1 && result.get()[0].isTensor()) + << "slice rows=" << rows << " forward failed"; + const auto& output = result.get()[0].toTensor(); + ASSERT_EQ(output.dim(), 2); + ASSERT_EQ(output.size(0), rows); + ASSERT_EQ(output.size(1), kSliceOutputColumns); + const size_t output_numel = static_cast(rows) * kSliceOutputColumns; + ASSERT_EQ(static_cast(output.numel()), output_numel); + const float* data = output.const_data_ptr(); + for (size_t i = 0; i < output_numel; ++i) { + const size_t row = i / kSliceOutputColumns; + const size_t column = i % kSliceOutputColumns; + const size_t input_index = row * kSliceInputColumns + column + 1; + ASSERT_EQ(data[i], slice_value(input_index, salt)) + << "slice rows=" << rows << " output index=" << i; + } +#ifdef WGPU_BACKEND_ENABLE_PROFILING + if (slice_profile_available()) { + expect_slice_profile(rows); + } +#endif +} + +void check_slice_dual_store() { + auto input = read_bin(g_dir + "/slice_dual_store.input.bin"); + auto golden0 = read_bin(g_dir + "/slice_dual_store.out0.golden.bin"); + auto golden1 = read_bin(g_dir + "/slice_dual_store.out1.golden.bin"); + const size_t input_numel = + static_cast(kSliceDualRows) * kSliceDualInputColumns; + const size_t output_numel = + static_cast(kSliceDualRows) * kSliceDualOutputColumns; + ASSERT_EQ(input.size(), input_numel); + ASSERT_EQ(golden0.size(), output_numel); + ASSERT_EQ(golden1.size(), output_numel); + + Module module(g_dir + "/slice_dual_store.pte"); + ASSERT_EQ(module.load_forward(), Error::Ok) << "load slice_dual_store.pte"; + auto tensor = make_tensor_ptr( + {kSliceDualRows, kSliceDualInputColumns}, std::move(input)); + auto result = module.forward({EValue(tensor)}); + ASSERT_TRUE(result.ok()) << "slice dual-store forward failed"; + ASSERT_EQ(result.get().size(), 2u); + const std::array*, 2> goldens = {&golden0, &golden1}; + for (size_t output_index = 0; output_index < goldens.size(); ++output_index) { + ASSERT_TRUE(result.get()[output_index].isTensor()); + const auto& output = result.get()[output_index].toTensor(); + ASSERT_EQ(output.dim(), 2); + ASSERT_EQ(output.size(0), kSliceDualRows); + ASSERT_EQ(output.size(1), kSliceDualOutputColumns); + ASSERT_EQ(static_cast(output.numel()), output_numel); + std::vector actual( + output.const_data_ptr(), + output.const_data_ptr() + output_numel); + EXPECT_EQ(max_err(actual, *goldens[output_index]), 0.0f) + << "slice dual-store output " << output_index; + } + +#ifdef WGPU_BACKEND_ENABLE_PROFILING + if (slice_profile_available()) { + const auto* context = get_default_webgpu_context(); + const auto& profile = context->querypool->results(); + const auto slice_count = + std::count_if(profile.begin(), profile.end(), [](const auto& duration) { + return duration.kernel_name == "slice"; + }); + EXPECT_EQ(slice_count, 1) << "dual-store must replace two slices with one"; + const auto slice = + std::find_if(profile.begin(), profile.end(), [](const auto& duration) { + return duration.kernel_name == "slice"; + }); + ASSERT_NE(slice, profile.end()); + EXPECT_EQ(slice->global_wg, (std::array{268, 268, 1})); + } +#endif +} + +float cat_value(size_t linear_index, int salt) { + return static_cast( + (linear_index * 19 + static_cast(salt) * 103) % kCatPatternSize); +} + +void check_cat_2d(Module& module, int rows, int salt) { + const size_t main_numel = static_cast(rows) * kCatInputColumns; + std::vector main_input(main_numel); + std::vector suffix(static_cast(rows)); + for (size_t i = 0; i < main_input.size(); ++i) { + main_input[i] = cat_value(i, salt); + } + for (size_t i = 0; i < suffix.size(); ++i) { + suffix[i] = cat_value(i, salt + 100); + } + + auto main_tensor = + make_tensor_ptr({rows, kCatInputColumns}, std::move(main_input)); + auto suffix_tensor = make_tensor_ptr({rows, 1}, std::move(suffix)); + auto result = module.forward({EValue(main_tensor), EValue(suffix_tensor)}); + ASSERT_TRUE( + result.ok() && result.get().size() == 1 && result.get()[0].isTensor()) + << "cat rows=" << rows << " forward failed"; + const auto& output = result.get()[0].toTensor(); + ASSERT_EQ(output.dim(), 2); + ASSERT_EQ(output.size(0), rows); + ASSERT_EQ(output.size(1), kCatOutputColumns); + const size_t output_numel = static_cast(rows) * kCatOutputColumns; + ASSERT_EQ(static_cast(output.numel()), output_numel); + const float* data = output.const_data_ptr(); + for (size_t i = 0; i < output_numel; ++i) { + const size_t row = i / kCatOutputColumns; + const size_t column = i % kCatOutputColumns; + const float expected = column < kCatInputColumns + ? cat_value(row * kCatInputColumns + column, salt) + : cat_value(row, salt + 100); + ASSERT_EQ(data[i], expected) + << "cat rows=" << rows << " output index=" << i; + } +} + // Dynamic quantized linear: input [M, kLinK] -> output [M, n]. kLinN is the // register-tiled/bicol config; kLinNShmem (N>=2048) routes to the shmem GEMM. constexpr int kLinK = 64; @@ -986,6 +1158,55 @@ TEST(DynamicShape, GeluCrosses2dDispatchBoundary) { } } +TEST(DynamicShape, SliceCrosses2dDispatchBoundary) { + if (std::getenv("WEBGPU_TEST_HEAVY") == nullptr) { + GTEST_SKIP() << "WEBGPU_TEST_HEAVY not set"; + } + const auto folded = + utils::fold_workgroup_count_2d(kSlice2dRows, 65535, "slice(test)"); + EXPECT_EQ(folded.x, 256); + EXPECT_EQ(folded.y, 256); + const auto linear = + utils::fold_workgroup_count_2d(kSlice1dRows, 65535, "slice(test)"); + EXPECT_EQ(linear.x, 65535); + EXPECT_EQ(linear.y, 1); + + Module module(g_dir + "/dyn_slice_2d.pte"); + ASSERT_EQ(module.load_forward(), Error::Ok) << "load dyn_slice_2d.pte"; + int salt = 1; + for (int rows : + {kSlice2dRows, kSlice1dRows, kSlice2dRows, kSlice1dRows, kSlice2dRows}) { + check_slice_2d(module, rows, salt++); + } +} + +TEST(DynamicShape, CatCrosses2dDispatchBoundary) { + if (std::getenv("WEBGPU_TEST_HEAVY") == nullptr) { + GTEST_SKIP() << "WEBGPU_TEST_HEAVY not set"; + } + const auto folded = utils::fold_workgroup_count_2d(65536, 65535, "cat(test)"); + EXPECT_EQ(folded.x, 256); + EXPECT_EQ(folded.y, 256); + const auto linear = utils::fold_workgroup_count_2d(65535, 65535, "cat(test)"); + EXPECT_EQ(linear.x, 65535); + EXPECT_EQ(linear.y, 1); + + Module module(g_dir + "/dyn_cat_2d.pte"); + ASSERT_EQ(module.load_forward(), Error::Ok) << "load dyn_cat_2d.pte"; + int salt = 1; + for (int rows : + {kCat2dRows, kCat1dRows, kCat2dRows, kCat1dRows, kCat2dRows}) { + check_cat_2d(module, rows, salt++); + } +} + +TEST(DynamicShape, SliceDualStoreWritesBothDestinations) { + if (std::getenv("WEBGPU_TEST_HEAVY") == nullptr) { + GTEST_SKIP() << "WEBGPU_TEST_HEAVY not set"; + } + check_slice_dual_store(); +} + TEST(DynamicShape, ExpandCopyRejectsDynamicShapesAtLoad) { const std::string path = g_dir + "/dyn_expand_copy.pte"; ASSERT_TRUE(std::ifstream(path).good()) << "missing dyn_expand_copy.pte"; @@ -2103,9 +2324,11 @@ int main(int argc, char** argv) { ctx = create_webgpu_context(); } catch (const std::exception& e) { std::printf("SKIP: no WebGPU device (%s)\n", e.what()); - return 0; + return required_device_failure_exit_code( + std::getenv("WEBGPU_REQUIRE_DEVICE") != nullptr); } set_default_webgpu_context(&ctx); + std::printf("WebGPU device acquired (native)\n"); const int rc = RUN_ALL_TESTS(); set_default_webgpu_context(nullptr); diff --git a/backends/webgpu/test/native/test_slice_chain.cpp b/backends/webgpu/test/native/test_slice_chain.cpp new file mode 100644 index 00000000000..844be0e6497 --- /dev/null +++ b/backends/webgpu/test/native/test_slice_chain.cpp @@ -0,0 +1,122 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include + +#include + +using executorch::backends::webgpu::WebGPUGraph; + +namespace { + +WebGPUGraph::SliceChain armed_chain() { + WebGPUGraph::SliceChain chain; + chain.valid = true; + chain.out_id = 7; + chain.dispatch_idx = 3; + chain.in_buffer = reinterpret_cast(0x1000); + chain.in_nbytes = 64; + chain.out_buffer = reinterpret_cast(0x2000); + chain.out_nbytes = 64; + chain.out_meta_buf = reinterpret_cast(0x3000); + chain.in_meta_buf = reinterpret_cast(0x4000); + chain.params_buf = reinterpret_cast(0x5000); + return chain; +} + +// A fresh graph must never observe another graph's offered chain. Before the +// chain became graph-instance state it lived in a function-local `static`, so +// two graphs shared one dispatch index and one set of buffer handles. +TEST(SliceChainLifecycle, SeparateGraphsDoNotShareState) { + WebGPUGraph first; + WebGPUGraph second; + + first.offer_slice_chain(armed_chain()); + + EXPECT_TRUE(first.slice_chain().valid); + EXPECT_FALSE(second.slice_chain().valid); + EXPECT_EQ(second.slice_chain().out_id, -1); + EXPECT_EQ(second.slice_chain().in_buffer, nullptr); + EXPECT_EQ(second.slice_chain().out_buffer, nullptr); +} + +// Interleaved offers must stay independent: arming the second graph must not +// disturb the first, and clearing one must not clear the other. +TEST(SliceChainLifecycle, InterleavedGraphsAreIndependent) { + WebGPUGraph first; + WebGPUGraph second; + + first.offer_slice_chain(armed_chain()); + WebGPUGraph::SliceChain other = armed_chain(); + other.out_id = 11; + other.dispatch_idx = 5; + second.offer_slice_chain(other); + + EXPECT_EQ(first.slice_chain().out_id, 7); + EXPECT_EQ(second.slice_chain().out_id, 11); + + first.clear_slice_chain(); + EXPECT_FALSE(first.slice_chain().valid); + EXPECT_TRUE(second.slice_chain().valid); + EXPECT_EQ(second.slice_chain().dispatch_idx, 5u); +} + +// A build that fails before completing must not leave a chain armed for the +// next build on the same graph. build() clears on entry and again via its +// scope guard; with no WebGPU context it fails early, which is the cheapest +// reachable failure path. +TEST(SliceChainLifecycle, FailedBuildClearsState) { + WebGPUGraph graph; + graph.offer_slice_chain(armed_chain()); + ASSERT_TRUE(graph.slice_chain().valid); + + try { + graph.build(nullptr, nullptr, 0, nullptr, {}); + } catch (...) { + // A build failure is expected here; the contract under test is the state. + } + + EXPECT_FALSE(graph.slice_chain().valid); + EXPECT_EQ(graph.slice_chain().out_id, -1); + EXPECT_EQ(graph.slice_chain().dispatch_idx, 0u); +} + +// Clearing must null every handle, so a later slice cannot re-bind a dispatch +// against buffers owned by a previous build. +TEST(SliceChainLifecycle, ClearDropsEveryStaleHandle) { + WebGPUGraph graph; + graph.offer_slice_chain(armed_chain()); + graph.clear_slice_chain(); + + const WebGPUGraph::SliceChain& chain = graph.slice_chain(); + EXPECT_FALSE(chain.valid); + EXPECT_EQ(chain.in_buffer, nullptr); + EXPECT_EQ(chain.out_buffer, nullptr); + EXPECT_EQ(chain.out_meta_buf, nullptr); + EXPECT_EQ(chain.in_meta_buf, nullptr); + EXPECT_EQ(chain.params_buf, nullptr); + EXPECT_EQ(chain.in_nbytes, 0u); + EXPECT_EQ(chain.out_nbytes, 0u); +} + +// Destroying a graph must not leave anything for the next graph to claim: a +// newly constructed graph always starts disarmed. +TEST(SliceChainLifecycle, DestructionLeavesNoResidueForTheNextGraph) { + { + WebGPUGraph doomed; + doomed.offer_slice_chain(armed_chain()); + ASSERT_TRUE(doomed.slice_chain().valid); + } + + WebGPUGraph fresh; + EXPECT_FALSE(fresh.slice_chain().valid); + EXPECT_EQ(fresh.slice_chain().out_id, -1); + EXPECT_EQ(fresh.slice_chain().params_buf, nullptr); +} + +} // namespace diff --git a/backends/webgpu/test/native/test_update_cache.cpp b/backends/webgpu/test/native/test_update_cache.cpp index dad859af669..d527ed3c1ac 100644 --- a/backends/webgpu/test/native/test_update_cache.cpp +++ b/backends/webgpu/test/native/test_update_cache.cpp @@ -6,16 +6,21 @@ * LICENSE file in the root directory of this source tree. */ +#include #include +#include #include #include #include +#include #include +#include #include #include #include +#include #include #include #include @@ -80,19 +85,15 @@ void run_case(const UpdateCacheCase& tc) { const auto& outputs = result.get(); ASSERT_TRUE(!outputs.empty() && outputs[0].isTensor()) << "no tensor output"; const auto& out_tensor = outputs[0].toTensor(); - ASSERT_EQ(static_cast(out_tensor.numel()), cnumel) - << "output numel " << (size_t)out_tensor.numel() << " != expected " - << cnumel; + ASSERT_EQ(out_tensor.dim(), 4); + ASSERT_EQ(out_tensor.size(0), 1); + ASSERT_EQ(out_tensor.size(1), tc.cmax); + ASSERT_EQ(out_tensor.size(2), tc.h); + ASSERT_EQ(out_tensor.size(3), tc.d); + ASSERT_EQ(static_cast(out_tensor.numel()), cnumel); const float* out_data = out_tensor.const_data_ptr(); - - float max_abs_err = 0.0f; - for (int i = 0; i < cnumel; i++) { - max_abs_err = std::max(max_abs_err, std::abs(out_data[i] - ref[i])); - } - // update_cache is a pure scatter copy: the output must be bit-exact. - EXPECT_EQ(max_abs_err, 0.0f) - << "update_cache[" << tc.name << "] not bit-exact (max abs error " - << max_abs_err << ", checked " << cnumel << " elements)"; + EXPECT_EQ(std::memcmp(out_data, ref.data(), ref.size() * sizeof(float)), 0) + << "update_cache[" << tc.name << "] not bit-exact"; } struct ReplayCase { @@ -138,11 +139,16 @@ void run_replay(const ReplayCase& rc) { ASSERT_TRUE(result.ok()) << "forward failed step " << step << " (error " << (int)result.error() << ")"; const auto& outputs = result.get(); - ASSERT_TRUE( - !outputs.empty() && outputs[0].isTensor() && - static_cast(outputs[0].toTensor().numel()) == cnumel) - << "bad cache output at step " << step; - const float* out_data = outputs[0].toTensor().const_data_ptr(); + ASSERT_EQ(outputs.size(), 1u) << "bad output count at step " << step; + ASSERT_TRUE(outputs[0].isTensor()) << "non-tensor output at step " << step; + const auto& output = outputs[0].toTensor(); + ASSERT_EQ(output.dim(), 4); + ASSERT_EQ(output.size(0), 1); + ASSERT_EQ(output.size(1), cmax); + ASSERT_EQ(output.size(2), rc.h); + ASSERT_EQ(output.size(3), rc.d); + ASSERT_EQ(static_cast(output.numel()), cnumel); + const float* out_data = output.const_data_ptr(); const int dst_offset = input_pos * rc.h * rc.d; for (int i = 0; i < vnumel; i++) { @@ -151,19 +157,233 @@ void run_replay(const ReplayCase& rc) { } } - float max_abs_err = 0.0f; - for (int i = 0; i < cnumel; i++) { - max_abs_err = std::max(max_abs_err, std::abs(out_data[i] - ref[i])); - cache[i] = out_data[i]; // thread the accumulated cache into the next step - } - // pure scatter copy: must be bit-exact - EXPECT_EQ(max_abs_err, 0.0f) - << "step " << step << " (S=" << s << ",pos=" << input_pos - << "): max abs error " << max_abs_err; + EXPECT_EQ(std::memcmp(out_data, ref.data(), ref.size() * sizeof(float)), 0) + << "step " << step << " (S=" << s << ",pos=" << input_pos << ")"; + std::memcpy(cache.data(), out_data, cache.size() * sizeof(float)); input_pos += s; } } +constexpr int kDynamicHeads = 2; +constexpr int kDynamicHeadDim = 4; +constexpr int kDynamicMaxCache = 1024; +constexpr int kUpdateCacheWorkgroupSize = 64; + +nlohmann::json current_attestation() { + return nlohmann::json::parse(webgpu_backend_execution_attestation_json()); +} + +const nlohmann::json* unique_update_cache_command( + const nlohmann::json& attestation) { + const auto& commands = attestation.at("canonicalCommands").at("commands"); + const nlohmann::json* match = nullptr; + for (const auto& command : commands) { + if (command.at("kind") == "compute" && command.at("enabled") == true && + command.at("identity") == "update_cache") { + EXPECT_EQ(match, nullptr) << "multiple enabled update_cache commands"; + match = &command; + } + } + EXPECT_NE(match, nullptr) << "missing enabled update_cache command"; + return match; +} + +void run_dynamic_success( + Module& module, + std::vector& full_cache, + int sequence, + int capacity, + int input_pos, + int salt, + uint64_t* last_ordinal) { + const int stride = kDynamicHeads * kDynamicHeadDim; + const int value_numel = sequence * stride; + const int cache_numel = capacity * stride; + std::vector value(static_cast(value_numel)); + for (int i = 0; i < value_numel; ++i) { + value[i] = static_cast(salt * 100000 + i); + } + std::vector cache( + full_cache.begin(), full_cache.begin() + cache_numel); + std::vector expected(cache); + std::memcpy( + expected.data() + static_cast(input_pos) * stride, + value.data(), + value.size() * sizeof(float)); + + auto value_tensor = make_tensor_ptr( + {1, sequence, kDynamicHeads, kDynamicHeadDim}, std::move(value)); + auto cache_tensor = make_tensor_ptr( + {1, capacity, kDynamicHeads, kDynamicHeadDim}, std::move(cache)); + auto position_tensor = + make_tensor_ptr({1}, std::vector{input_pos}); + auto result = module.forward( + {EValue(value_tensor), EValue(cache_tensor), EValue(position_tensor)}); + ASSERT_TRUE(result.ok()) << "S=" << sequence << " C=" << capacity + << " pos=" << input_pos + << " error=" << static_cast(result.error()); + size_t cache_output_count = 0; + for (const auto& output_value : result.get()) { + if (!output_value.isTensor()) { + continue; + } + const auto& output = output_value.toTensor(); + if (output.dim() != 4 || output.size(0) != 1 || + output.size(1) != capacity || output.size(2) != kDynamicHeads || + output.size(3) != kDynamicHeadDim) { + continue; + } + ++cache_output_count; + ASSERT_EQ(static_cast(output.numel()), cache_numel); + const float* output_data = output.const_data_ptr(); + ASSERT_EQ( + std::memcmp( + output_data, expected.data(), expected.size() * sizeof(float)), + 0); + std::memcpy( + full_cache.data(), output_data, expected.size() * sizeof(float)); + } + ASSERT_EQ(cache_output_count, 1u); + + const nlohmann::json attestation = current_attestation(); + ASSERT_TRUE(attestation.at("completed").get()); + const uint64_t ordinal = attestation.at("executionOrdinal").get(); + if (*last_ordinal != 0) { + EXPECT_EQ(ordinal, *last_ordinal + 1); + } + *last_ordinal = ordinal; + const auto* command = unique_update_cache_command(attestation); + ASSERT_NE(command, nullptr); + const uint32_t expected_grid = static_cast( + (value_numel + kUpdateCacheWorkgroupSize - 1) / + kUpdateCacheWorkgroupSize); + EXPECT_EQ( + command->at("grid"), + nlohmann::json::array({expected_grid, 1u, 1u})); +} + +void expect_dynamic_failure( + Module& module, + const std::vector& full_cache, + int sequence, + int capacity, + int input_pos) { + const int stride = kDynamicHeads * kDynamicHeadDim; + const std::string before = webgpu_backend_execution_attestation_json(); + for (int attempt = 0; attempt < 2; ++attempt) { + std::vector value( + static_cast(sequence * stride), 7.0f); + std::vector cache( + full_cache.begin(), full_cache.begin() + capacity * stride); + const std::vector expected_value(value); + const std::vector expected_cache(cache); + auto value_tensor = make_tensor_ptr( + {1, sequence, kDynamicHeads, kDynamicHeadDim}, std::move(value)); + auto cache_tensor = make_tensor_ptr( + {1, capacity, kDynamicHeads, kDynamicHeadDim}, std::move(cache)); + auto position_tensor = + make_tensor_ptr({1}, std::vector{input_pos}); + + auto result = module.forward( + {EValue(value_tensor), EValue(cache_tensor), EValue(position_tensor)}); + + ASSERT_FALSE(result.ok()); + EXPECT_EQ(result.error(), Error::Internal); + EXPECT_EQ( + std::memcmp( + value_tensor->const_data_ptr(), + expected_value.data(), + expected_value.size() * sizeof(float)), + 0); + EXPECT_EQ( + std::memcmp( + cache_tensor->const_data_ptr(), + expected_cache.data(), + expected_cache.size() * sizeof(float)), + 0); + EXPECT_EQ(webgpu_backend_execution_attestation_json(), before); + } +} + +void run_intermediate_success( + Module& module, + std::vector& full_cache, + int sequence, + int input_pos, + int salt, + uint64_t* last_ordinal) { + constexpr int kCapacity = 768; + const int stride = kDynamicHeads * kDynamicHeadDim; + const int value_numel = sequence * stride; + const int cache_numel = kCapacity * stride; + std::vector value(static_cast(value_numel)); + std::vector transformed(static_cast(value_numel)); + for (int i = 0; i < value_numel; ++i) { + value[i] = static_cast((i + salt) % 33 - 16) * 0.125f; + transformed[i] = 1.0f / (1.0f + std::exp(-value[i])); + } + std::vector cache(full_cache.begin(), full_cache.begin() + cache_numel); + const std::vector before(cache); + + auto value_tensor = make_tensor_ptr( + {1, sequence, kDynamicHeads, kDynamicHeadDim}, std::move(value)); + auto cache_tensor = make_tensor_ptr( + {1, kCapacity, kDynamicHeads, kDynamicHeadDim}, std::move(cache)); + auto position_tensor = + make_tensor_ptr({1}, std::vector{input_pos}); + auto result = module.forward( + {EValue(value_tensor), EValue(cache_tensor), EValue(position_tensor)}); + ASSERT_TRUE(result.ok()) << "S=" << sequence << " pos=" << input_pos + << " error=" << static_cast(result.error()); + + size_t cache_output_count = 0; + for (const auto& output_value : result.get()) { + if (!output_value.isTensor()) { + continue; + } + const auto& output = output_value.toTensor(); + if (output.dim() != 4 || output.size(0) != 1 || + output.size(1) != kCapacity || output.size(2) != kDynamicHeads || + output.size(3) != kDynamicHeadDim) { + continue; + } + ++cache_output_count; + const float* output_data = output.const_data_ptr(); + const size_t write_begin = static_cast(input_pos) * stride; + const size_t write_end = write_begin + transformed.size(); + for (size_t i = 0; i < static_cast(cache_numel); ++i) { + if (i >= write_begin && i < write_end) { + EXPECT_NEAR(output_data[i], transformed[i - write_begin], 1.0e-3f); + } else { + EXPECT_EQ( + std::memcmp(output_data + i, before.data() + i, sizeof(float)), + 0); + } + } + std::memcpy( + full_cache.data(), + output_data, + static_cast(cache_numel) * sizeof(float)); + } + ASSERT_EQ(cache_output_count, 1u); + + const nlohmann::json attestation = current_attestation(); + ASSERT_TRUE(attestation.at("completed").get()); + const uint64_t ordinal = attestation.at("executionOrdinal").get(); + if (*last_ordinal != 0) { + EXPECT_EQ(ordinal, *last_ordinal + 1); + } + *last_ordinal = ordinal; + const auto* command = unique_update_cache_command(attestation); + ASSERT_NE(command, nullptr); + const uint32_t expected_grid = static_cast( + (value_numel + kUpdateCacheWorkgroupSize - 1) / + kUpdateCacheWorkgroupSize); + EXPECT_EQ( + command->at("grid"), + nlohmann::json::array({expected_grid, 1u, 1u})); +} + struct NegativeCase { const char* name; const char* guard; @@ -214,6 +434,48 @@ TEST(UpdateCache, Negative) { } } +TEST(UpdateCache, DynamicSymIntShapeCapacityBoundsAndRecovery) { + Module module(g_dir + "/dynamic.pte"); + ASSERT_EQ(module.load_forward(), Error::Ok) << "could not load dynamic.pte"; + std::vector full_cache( + static_cast(kDynamicMaxCache) * kDynamicHeads * kDynamicHeadDim); + for (size_t i = 0; i < full_cache.size(); ++i) { + full_cache[i] = static_cast(i) + 0.25f; + } + uint64_t ordinal = 0; + + run_dynamic_success(module, full_cache, 512, 1024, 0, 1, &ordinal); + run_dynamic_success(module, full_cache, 1, 1024, 0, 2, &ordinal); + run_dynamic_success(module, full_cache, 1, 768, 0, 3, &ordinal); + run_dynamic_success(module, full_cache, 1, 768, 512, 4, &ordinal); + expect_dynamic_failure(module, full_cache, 1, 512, 512); + run_dynamic_success(module, full_cache, 1, 768, 513, 5, &ordinal); + run_dynamic_success(module, full_cache, 512, 1024, 0, 6, &ordinal); + run_dynamic_success(module, full_cache, 1, 1024, 0, 7, &ordinal); + + expect_dynamic_failure(module, full_cache, 1, 1024, -1); + run_dynamic_success(module, full_cache, 1, 1024, 1, 8, &ordinal); + expect_dynamic_failure(module, full_cache, 512, 1024, 513); + run_dynamic_success(module, full_cache, 1, 1024, 2, 9, &ordinal); +} + +TEST(UpdateCache, DynamicIntermediateValueRefreshesAfterTensorFixpoint) { + Module module(g_dir + "/dynamic_intermediate.pte"); + ASSERT_EQ(module.load_forward(), Error::Ok) + << "could not load dynamic_intermediate.pte"; + std::vector full_cache( + static_cast(768) * kDynamicHeads * kDynamicHeadDim); + for (size_t i = 0; i < full_cache.size(); ++i) { + full_cache[i] = static_cast(i) + 0.75f; + } + uint64_t ordinal = 0; + + run_intermediate_success(module, full_cache, 512, 0, 1, &ordinal); + run_intermediate_success(module, full_cache, 1, 0, 2, &ordinal); + run_intermediate_success(module, full_cache, 512, 0, 3, &ordinal); + run_intermediate_success(module, full_cache, 1, 512, 4, &ordinal); +} + int main(int argc, char** argv) { ::testing::InitGoogleTest(&argc, argv); @@ -232,10 +494,12 @@ int main(int argc, char** argv) { try { ctx = create_webgpu_context(); } catch (const std::exception& e) { - std::printf("SKIP: %s\n", e.what()); - return 0; + std::printf("SKIP: no WebGPU device (%s)\n", e.what()); + return required_device_failure_exit_code( + std::getenv("WEBGPU_REQUIRE_DEVICE") != nullptr); } set_default_webgpu_context(&ctx); + std::printf("WebGPU device acquired (native)\n"); const int rc = RUN_ALL_TESTS(); set_default_webgpu_context(nullptr); diff --git a/backends/webgpu/test/native/test_update_cache_state.cpp b/backends/webgpu/test/native/test_update_cache_state.cpp new file mode 100644 index 00000000000..051a9a021fc --- /dev/null +++ b/backends/webgpu/test/native/test_update_cache_state.cpp @@ -0,0 +1,125 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include + +#include + +#include +#include +#include +#include + +using namespace executorch::backends::webgpu; + +namespace { + +LiveUpdateCacheInputs valid_inputs() { + return { + {1, 512, 2, 4}, + {1, 1024, 2, 4}, + 4, + 4, + 0, + 64, + 65535, + }; +} + +TEST(UpdateCacheState, ComputesCompleteStateBeforeOneCommit) { + int commits = 0; + LiveUpdateCacheState committed = {}; + + refresh_live_update_cache_state( + valid_inputs(), [&](const LiveUpdateCacheState& state) { + ++commits; + committed = state; + }); + + EXPECT_EQ(commits, 1); + EXPECT_EQ(committed.params.numel, 4096u); + EXPECT_EQ(committed.params.dst_offset, 0u); + EXPECT_EQ(committed.params.cache_numel, 8192u); + EXPECT_EQ(committed.workgroup_count_x, 64u); +} + +TEST(UpdateCacheState, RejectsEveryInvalidInputBeforeCommit) { + std::vector invalid; + auto input = valid_inputs(); + input.value_dims = {1, 1, 2}; + invalid.push_back(input); + input = valid_inputs(); + input.value_dims[1] = 0; + invalid.push_back(input); + input = valid_inputs(); + input.value_dims[0] = 2; + invalid.push_back(input); + input = valid_inputs(); + input.cache_dims[2] = 3; + invalid.push_back(input); + input = valid_inputs(); + input.start_pos = -1; + invalid.push_back(input); + input = valid_inputs(); + input.start_pos = 513; + invalid.push_back(input); + input = valid_inputs(); + input.max_workgroups_per_dimension = 63; + invalid.push_back(input); + input = valid_inputs(); + input.workgroup_size = 0; + invalid.push_back(input); + input = valid_inputs(); + input.value_dims = {1, 1, 65536, 65536}; + input.cache_dims = input.value_dims; + invalid.push_back(input); + + for (const auto& candidate : invalid) { + int commits = 0; + EXPECT_THROW( + refresh_live_update_cache_state( + candidate, [&](const LiveUpdateCacheState&) { ++commits; }), + std::runtime_error); + EXPECT_EQ(commits, 0); + } +} + +TEST(UpdateCacheState, RejectsStartPositionMultiplicationOverflow) { + LiveUpdateCacheInputs input = { + {1, 1, 1, 3}, + {1, 2, 1, 3}, + 4, + 4, + INT64_C(6148914691236517206), + 64, + 65535, + }; + int commits = 0; + + EXPECT_THROW( + refresh_live_update_cache_state( + input, [&](const LiveUpdateCacheState&) { ++commits; }), + std::runtime_error); + EXPECT_EQ(commits, 0); +} + +TEST(UpdateCacheState, RejectsNumelOverflowBeforeCommit) { + auto input = valid_inputs(); + input.value_dims = { + 1, std::numeric_limits::max(), 2, 2}; + input.cache_dims = input.value_dims; + int commits = 0; + + EXPECT_THROW( + refresh_live_update_cache_state( + input, [&](const LiveUpdateCacheState&) { ++commits; }), + std::runtime_error); + EXPECT_EQ(commits, 0); +} + +} // namespace diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index 6a019ab8e9f..a96a96ab56f 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -1051,7 +1051,16 @@ def _unsqueeze_suite() -> WebGPUTestSuite: @register_op_test("slice") def _slice_suite() -> WebGPUTestSuite: - return _fn_config_suite(SliceModule, _SLICE_CONFIGS) + suite = _fn_config_suite(SliceModule, _SLICE_CONFIGS) + suite.cases.append( + Case( + name="folded_2d_65536_workgroups", + construct={"fn": lambda x: x[:, 1:65]}, + inputs=((65536, 65),), + heavy=True, + ) + ) + return suite @register_op_test("permute") @@ -1070,7 +1079,7 @@ def _permute_suite() -> WebGPUTestSuite: @register_op_test("cat") def _cat_suite() -> WebGPUTestSuite: # CONFIGS: name -> (list_of_input_shapes, dim). Variadic input count per case. - return WebGPUTestSuite( + suite = WebGPUTestSuite( module_factory=lambda dim: CatModule(dim), cases=[ Case(name=n, construct={"dim": dim}, inputs=tuple(shapes)) @@ -1078,6 +1087,15 @@ def _cat_suite() -> WebGPUTestSuite: ], golden_dtype="float32", # concatenation copies values; fp64 bit-identical ) + suite.cases.append( + Case( + name="folded_2d_full_output", + construct={"dim": 1}, + inputs=((65536, 65), (65536, 1)), + heavy=True, + ) + ) + return suite from executorch.backends.webgpu.test.ops.test_gelu import ( diff --git a/backends/webgpu/test/op_tests/op_test_driver.cpp b/backends/webgpu/test/op_tests/op_test_driver.cpp index 93b2df9daf8..d9176c58253 100644 --- a/backends/webgpu/test/op_tests/op_test_driver.cpp +++ b/backends/webgpu/test/op_tests/op_test_driver.cpp @@ -10,6 +10,7 @@ // per manifest case, runs forward() on the device, and compares vs golden. #include +#include #include #include #include @@ -218,8 +219,10 @@ int main(int argc, char** argv) { ctx = create_webgpu_context(); } catch (const std::exception& ex) { std::printf("SKIP: no WebGPU device (%s)\n", ex.what()); - return 0; + return required_device_failure_exit_code( + std::getenv("WEBGPU_REQUIRE_DEVICE") != nullptr); } + std::printf("WebGPU device acquired (native)\n"); set_default_webgpu_context(&ctx); auto entries = parse_manifest(manifest); diff --git a/backends/webgpu/test/ops/dynamic_shape/test_dynamic_shape_export.py b/backends/webgpu/test/ops/dynamic_shape/test_dynamic_shape_export.py index d9f666622ee..1842cd436ef 100644 --- a/backends/webgpu/test/ops/dynamic_shape/test_dynamic_shape_export.py +++ b/backends/webgpu/test/ops/dynamic_shape/test_dynamic_shape_export.py @@ -13,6 +13,7 @@ goldens. """ +import dataclasses import math import os import unittest @@ -21,8 +22,10 @@ from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner from executorch.backends.webgpu.test.ops.test_conv1d_pw import Conv1dModule from executorch.backends.webgpu.test.ops.test_gelu import GeluModule -from executorch.exir import to_edge_transform_and_lower +from executorch.backends.webgpu.test.ops.test_slice import SliceModule +from executorch.exir import EdgeCompileConfig, to_edge, to_edge_transform_and_lower from executorch.exir.backend.utils import get_delegates, get_non_lowered_nodes +from executorch.exir.dialects._ops import ops as exir_ops MAXS = 128 # upper bound for the dynamic seq-len dim (within the 1D dispatch cap) HIDDEN = 64 @@ -197,6 +200,17 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return x.expand((4, -1)).clone() +class SliceDualStoreModule(torch.nn.Module): + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + first = torch.ops.aten.slice_copy.Tensor(x, 1, 1, 8961, 1) + return first, first + + +class DynamicCatModule(torch.nn.Module): + def forward(self, x: torch.Tensor, suffix: torch.Tensor) -> torch.Tensor: + return torch.cat((x, suffix), dim=1) + + def _ramp(shape) -> torch.Tensor: n = 1 for d in shape: @@ -312,6 +326,109 @@ def export_dynamic_gelu_boundary_cases(out_dir: str) -> None: ) +def export_dynamic_slice_boundary_case(out_dir: str) -> None: + """Write a reusable slice graph crossing the per-dimension grid cap.""" + max_rows = 65536 + model = SliceModule(lambda x: x[:, 1:65]).eval() + rows_dim = torch.export.Dim("slice_rows", min=1, max=max_rows) + _export( + model, + (torch.empty((max_rows, 65), dtype=torch.float32),), + {"x": {0: rows_dim}}, + os.path.join(out_dir, "dyn_slice_2d.pte"), + ) + + +def export_dynamic_cat_boundary_case(out_dir: str) -> None: + """Write a reusable Cat graph crossing the per-dimension grid cap.""" + max_rows = 65536 + model = DynamicCatModule().eval() + rows_dim = torch.export.Dim("cat_rows", min=1, max=max_rows) + _export( + model, + ( + torch.empty((max_rows, 65), dtype=torch.float32), + torch.empty((max_rows, 1), dtype=torch.float32), + ), + {"x": {0: rows_dim}, "suffix": {0: rows_dim}}, + os.path.join(out_dir, "dyn_cat_2d.pte"), + ) + + +def export_slice_dual_store_case(out_dir: str) -> None: + """Write the retained two-slice edge graph used by the dual-store route.""" + rows = 512 + input_columns = 8961 + output_columns = 8960 + model = SliceDualStoreModule().eval() + input_tensor = ( + torch.arange(rows * input_columns, dtype=torch.float32) + .remainder(4093) + .reshape(rows, input_columns) + ) + exported = torch.export.export(model, (input_tensor,), strict=True) + edge = to_edge( + exported, + compile_config=EdgeCompileConfig(_check_ir_validity=False), + ) + edge_program = edge.exported_program() + graph = edge_program.graph_module.graph + first = next( + node + for node in graph.nodes + if node.target == exir_ops.edge.aten.slice_copy.Tensor + ) + output = next(node for node in graph.nodes if node.op == "output") + with graph.inserting_after(first): + second = graph.call_function( + exir_ops.edge.aten.slice_copy.Tensor, + args=(first, 1, 0, output_columns, 1), + ) + second.meta = first.meta.copy() + output.args = ((first, second),) + edge_program.graph_signature.output_specs[1] = dataclasses.replace( + edge_program.graph_signature.output_specs[1], + arg=dataclasses.replace( + edge_program.graph_signature.output_specs[1].arg, + name=second.name, + ), + ) + edge_program.graph_module.recompile() + if ( + sum(node.target == exir_ops.edge.aten.slice_copy.Tensor for node in graph.nodes) + != 2 + ): + raise RuntimeError("slice dual-store fixture must retain two slice_copy ops") + lowered = edge.to_backend(VulkanPartitioner()) + lowered_graph = lowered.exported_program().graph_module.graph + delegates = get_delegates(lowered_graph) + portable = get_non_lowered_nodes(lowered_graph) + if len(delegates) != 1 or portable: + raise RuntimeError( + "slice dual-store fixture must be one fully delegated subgraph: " + f"delegates={len(delegates)}, portable={portable}" + ) + program = lowered.to_executorch() + with open(os.path.join(out_dir, "slice_dual_store.pte"), "wb") as file: + file.write(program.buffer) + with torch.no_grad(): + first_golden = torch.ops.aten.slice_copy.Tensor( + input_tensor, 1, 1, input_columns, 1 + ) + second_golden = torch.ops.aten.slice_copy.Tensor( + first_golden, 1, 0, output_columns, 1 + ) + input_tensor.numpy().astype(" None: """Write a dynamic expand_copy graph that the runtime must reject at load.""" model = DynamicExpandCopyModule().eval() @@ -337,6 +454,9 @@ def export_dynamic_shape_cases(out_dir: str) -> None: export_dynamic_expand_copy_rejection_case(out_dir) if os.environ.get("WEBGPU_TEST_HEAVY"): export_dynamic_gelu_boundary_cases(out_dir) + export_dynamic_slice_boundary_case(out_dir) + export_dynamic_cat_boundary_case(out_dir) + export_slice_dual_store_case(out_dir) s_dim = torch.export.Dim("s", min=1, max=MAXS) # 1) Single dynamic rms_norm, graph built at S=MAXS (upper bound). diff --git a/backends/webgpu/test/ops/test_et_vk_sdpa.py b/backends/webgpu/test/ops/test_et_vk_sdpa.py index c999fd01a1f..d6aaceb20e0 100644 --- a/backends/webgpu/test/ops/test_et_vk_sdpa.py +++ b/backends/webgpu/test/ops/test_et_vk_sdpa.py @@ -27,7 +27,7 @@ import torch import torch.nn.functional as F -from executorch.backends.vulkan import VulkanPartitioner +from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner from executorch.exir import to_edge_transform_and_lower NEG_INF = -1e30 @@ -102,6 +102,34 @@ def _lower(cfg: SdpaConfig, q, k, v): class TestEtVkSdpa(unittest.TestCase): + def test_negative_infinity_mask_matches_eager(self) -> None: + # Exercises the production op, not a local reimplementation: masked + # positions must contribute nothing and must not leak NaN. An exact + # -inf (0xff800000) is the value the QK elision keys on. + for cfg in CONFIGS: + with self.subTest(config=cfg.name): + q, k, v = _qkv(cfg) + scale = 1.0 / math.sqrt(cfg.d) + mask = torch.zeros(cfg.s_q, cfg.s_kv) + # Mask the strict upper triangle, leaving every row with at + # least one unmasked position so softmax stays well defined. + mask.masked_fill_( + torch.triu( + torch.ones(cfg.s_q, cfg.s_kv, dtype=torch.bool), diagonal=1 + ), + float("-inf"), + ) + self.assertEqual( + torch.tensor(float("-inf")).view(torch.int32).item(), + -0x800000, + ) + got = torch.ops.et_vk.sdpa.default(q, k, v, mask, scale) + ref = F.scaled_dot_product_attention( + q, k, v, attn_mask=mask, scale=scale + ) + self.assertFalse(torch.isnan(got).any()) + torch.testing.assert_close(got, ref, atol=1e-4, rtol=1e-3) + def test_export_delegates(self) -> None: for cfg in CONFIGS: with self.subTest(config=cfg.name): diff --git a/backends/webgpu/test/ops/test_rope_hf_single.py b/backends/webgpu/test/ops/test_rope_hf_single.py new file mode 100644 index 00000000000..15de88901c3 --- /dev/null +++ b/backends/webgpu/test/ops/test_rope_hf_single.py @@ -0,0 +1,44 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import unittest + +import executorch.backends.vulkan.custom_ops_lib # noqa: F401 +import torch + + +class RopeHfSingleTest(unittest.TestCase): + def test_reference_matches_rotate_half(self) -> None: + x = torch.tensor( + [[[[1.0, 2.0, 3.0, 4.0], [-1.0, -2.0, -3.0, -4.0]]]] + ) + freqs_cos = torch.tensor( + [[1.0, 1.0, 1.0, 1.0], [0.5, 0.25, 0.5, 0.25]] + ) + freqs_sin = torch.tensor( + [[0.0, 0.0, 0.0, 0.0], [0.5, 0.75, 0.5, 0.75]] + ) + expected = torch.tensor( + [[[[ + -1.0, + -2.5, + 2.0, + 2.5, + ], [ + 1.0, + 2.5, + -2.0, + -2.5, + ]]]] + ) + + actual = torch.ops.et_vk.apply_rotary_emb_hf_single.default( + x, freqs_cos, freqs_sin, 1 + ) + + torch.testing.assert_close(expected, actual) + + diff --git a/backends/webgpu/test/ops/test_update_cache.py b/backends/webgpu/test/ops/test_update_cache.py index a25321bdd7d..5824572dfcb 100644 --- a/backends/webgpu/test/ops/test_update_cache.py +++ b/backends/webgpu/test/ops/test_update_cache.py @@ -21,9 +21,20 @@ # Importing custom_ops registers torch.ops.llama.update_cache (the schema lives # in the C++ AOT lib loaded here). -from executorch.backends.vulkan import VulkanPartitioner +from executorch.backends.vulkan.partitioner.vulkan_partitioner import ( + VulkanPartitioner, +) +from executorch.backends.vulkan.serialization.vulkan_graph_serialize import ( + extract_vk_flatbuffer, + flatbuffer_to_vk_graph, +) +from executorch.examples.models.gemma4.webgpu_partitioner import ( + build_webgpu_partitioner, +) from executorch.exir import to_edge_transform_and_lower +from executorch.exir.schema import DataLocation, DelegateCall, KernelCall from executorch.extension.llm.custom_ops import custom_ops # noqa: F401 +from torch.export.graph_signature import InputKind, OutputKind class UpdateCacheModule(torch.nn.Module): @@ -37,7 +48,221 @@ def forward(self, value: torch.Tensor, cache: torch.Tensor) -> torch.Tensor: return torch.ops.llama.update_cache(value, cache, self.input_pos) +class DynamicUpdateCacheModule(torch.nn.Module): + """Writes at the live scalar selected from the position tensor.""" + + def forward( + self, + value: torch.Tensor, + cache: torch.Tensor, + input_pos: torch.Tensor, + ) -> torch.Tensor: + return torch.ops.llama.update_cache(value, cache, input_pos[0].item()) + + +class RegisteredBufferDynamicUpdateCacheModule(torch.nn.Module): + """Mirrors Gemma's non-persistent registered KV cache.""" + + def __init__(self) -> None: + super().__init__() + self.register_buffer( + "cache", + torch.zeros(1, 1024, 2, 4), + persistent=False, + ) + + def forward( + self, + value: torch.Tensor, + input_pos: torch.Tensor, + ) -> torch.Tensor: + torch.ops.llama.update_cache(value, self.cache, input_pos[0].item()) + return self.cache[:, :1] + + +class IntermediateDynamicUpdateCacheModule(torch.nn.Module): + """Feeds an intermediate value with a live sequence shape to the cache.""" + + def forward( + self, + value: torch.Tensor, + cache: torch.Tensor, + input_pos: torch.Tensor, + ) -> torch.Tensor: + produced = torch.sigmoid(value) + return torch.ops.llama.update_cache( + produced, + cache, + input_pos[0].item(), + ) + + +def _lower_dynamic_update_cache( + model: torch.nn.Module, + inputs: tuple[torch.Tensor, ...], + dynamic_shapes: tuple[object, ...], +) -> tuple[torch.export.ExportedProgram, object]: + with torch._dynamo.config.patch(capture_scalar_outputs=True): + exported = torch.export.export( + model, + inputs, + dynamic_shapes=dynamic_shapes, + ) + lowered = to_edge_transform_and_lower( + exported, + partitioner=[build_webgpu_partitioner("8da4w+emb4")], + ).to_executorch() + return exported, lowered + + +def _export_dynamic_update_cache() -> tuple[torch.export.ExportedProgram, object]: + value = torch.zeros(1, 512, 2, 4) + cache = torch.zeros(1, 1024, 2, 4) + input_pos = torch.zeros(1, dtype=torch.long) + value_seq = torch.export.Dim("value_seq", min=1, max=512) + cache_capacity = torch.export.Dim("cache_capacity", min=512, max=1024) + return _lower_dynamic_update_cache( + DynamicUpdateCacheModule(), + (value, cache, input_pos), + ({1: value_seq}, {1: cache_capacity}, None), + ) + + +def _export_registered_program() -> tuple[ + torch.export.ExportedProgram, + tuple[torch.Tensor, torch.Tensor], + tuple[object, object], +]: + value = torch.zeros(1, 512, 2, 4) + input_pos = torch.zeros(1, dtype=torch.long) + value_seq = torch.export.Dim("registered_value_seq", min=1, max=512) + inputs = (value, input_pos) + dynamic_shapes = ({1: value_seq}, None) + with torch._dynamo.config.patch(capture_scalar_outputs=True): + exported = torch.export.export( + RegisteredBufferDynamicUpdateCacheModule(), + inputs, + dynamic_shapes=dynamic_shapes, + ) + return exported, inputs, dynamic_shapes + + +def _export_registered_dynamic_update_cache() -> tuple[ + torch.export.ExportedProgram, object +]: + contract, inputs, dynamic_shapes = _export_registered_program() + _, lowered = _lower_dynamic_update_cache( + RegisteredBufferDynamicUpdateCacheModule(), inputs, dynamic_shapes + ) + return contract, lowered + + +def _export_intermediate_dynamic_update_cache() -> tuple[ + torch.export.ExportedProgram, object +]: + value = torch.zeros(1, 512, 2, 4) + cache = torch.zeros(1, 768, 2, 4) + input_pos = torch.zeros(1, dtype=torch.long) + value_seq = torch.export.Dim("intermediate_value_seq", min=1, max=512) + return _lower_dynamic_update_cache( + IntermediateDynamicUpdateCacheModule(), + (value, cache, input_pos), + ({1: value_seq}, None, None), + ) + + class TestUpdateCache(unittest.TestCase): + def _assert_live_update_cache_symint( + self, + exported: torch.export.ExportedProgram, + ) -> None: + update_nodes = [ + node + for node in exported.graph_module.graph.nodes + if node.target == torch.ops.llama.update_cache.default + ] + self.assertEqual(1, len(update_nodes)) + start_pos = update_nodes[0].args[2] + self.assertIsInstance(start_pos, torch.fx.Node) + self.assertIsInstance(start_pos.meta.get("val"), torch.SymInt) + + def _assert_exact_delegate_chain( + self, + program: object, + expected_chain: list[str] | None = None, + ) -> None: + self.assertEqual(1, len(program.execution_plan)) + plan = program.execution_plan[0] + self.assertEqual( + ["VulkanBackend"], + [delegate.id for delegate in plan.delegates], + ) + self.assertEqual(1, len(plan.delegates)) + delegate = plan.delegates[0] + self.assertEqual(DataLocation.INLINE, delegate.processed.location) + self.assertGreaterEqual(delegate.processed.index, 0) + self.assertLess(delegate.processed.index, len(program.backend_delegate_data)) + payload = program.backend_delegate_data[delegate.processed.index].data + vk_graph = flatbuffer_to_vk_graph(extract_vk_flatbuffer(payload)) + self.assertEqual( + expected_chain + or ["et_vk.select_as_symint.default", "update_cache.default"], + [operator.name for operator in vk_graph.chain], + ) + + def _assert_user_cache_writeback( + self, + program: object, + expected_delegate_chain: list[str], + ) -> None: + self.assertEqual(1, len(program.execution_plan)) + plan = program.execution_plan[0] + calls = [ + instruction.instr_args + for chain in plan.chains + for instruction in chain.instructions + ] + delegate_calls = [call for call in calls if isinstance(call, DelegateCall)] + kernel_calls = [call for call in calls if isinstance(call, KernelCall)] + self.assertEqual(1, len(delegate_calls), repr(plan)) + self.assertEqual(3, len(kernel_calls), repr(plan)) + self.assertEqual( + [ + ("aten::copy", "out"), + ("aten::copy", "out"), + ("aten::copy_", ""), + ], + [ + ( + plan.operators[call.op_index].name, + plan.operators[call.op_index].overload, + ) + for call in kernel_calls + ], + ) + self.assertEqual( + [("aten::copy", "out"), ("aten::copy_", "")], + [(operator.name, operator.overload) for operator in plan.operators], + ) + + cache_input = plan.inputs[1] + first_stage, second_stage, writeback = kernel_calls + delegate_output = first_stage.args[1] + self.assertIn(delegate_output, delegate_calls[0].args) + self.assertNotIn(delegate_output, plan.inputs) + self.assertEqual(cache_input, first_stage.args[0]) + self.assertEqual(first_stage.args[-2], first_stage.args[-1]) + first_stage_output = first_stage.args[-1] + self.assertEqual(cache_input, second_stage.args[0]) + self.assertEqual(first_stage_output, second_stage.args[1]) + self.assertEqual(second_stage.args[-2], second_stage.args[-1]) + second_stage_output = second_stage.args[-1] + self.assertEqual(cache_input, writeback.args[0]) + self.assertEqual(second_stage_output, writeback.args[1]) + self.assertEqual(cache_input, writeback.args[-1]) + self.assertIn(cache_input, plan.outputs) + self._assert_exact_delegate_chain(program, expected_delegate_chain) + def _export_and_check(self, model, example_inputs) -> None: ep = torch.export.export(model, example_inputs) et_program = to_edge_transform_and_lower( @@ -64,6 +289,79 @@ def test_update_cache_gqa_shapes(self) -> None: cache = torch.zeros(1, 16, 2, 8) self._export_and_check(UpdateCacheModule(0), (value, cache)) + def test_registered_cache_keeps_live_position_and_has_no_portable_call( + self, + ) -> None: + exported, lowered = _export_registered_dynamic_update_cache() + self._assert_live_update_cache_symint(exported) + buffers = [ + input_spec + for input_spec in exported.graph_signature.input_specs + if input_spec.kind == InputKind.BUFFER + ] + self.assertEqual(1, len(buffers)) + self.assertEqual("cache", buffers[0].target) + self.assertFalse(buffers[0].persistent) + self.assertTrue( + all( + output.target != "cache" + for output in exported.graph_signature.output_specs + if output.kind == OutputKind.USER_OUTPUT + ) + ) + + program = lowered.executorch_program + plan = program.execution_plan[0] + calls = [ + instruction.instr_args + for chain in plan.chains + for instruction in chain.instructions + ] + self.assertEqual([], plan.operators) + self.assertEqual([], [call for call in calls if isinstance(call, KernelCall)]) + self.assertEqual( + 1, + len([call for call in calls if isinstance(call, DelegateCall)]), + ) + self._assert_exact_delegate_chain( + program, + [ + "et_vk.prepack.default", + "et_vk.select_as_symint.default", + "update_cache.default", + "aten.slice_copy.Tensor", + ], + ) + + def test_user_cache_keeps_live_position_and_exact_mutation_writeback( + self, + ) -> None: + exported, lowered = _export_dynamic_update_cache() + self._assert_live_update_cache_symint(exported) + + self._assert_user_cache_writeback( + lowered.executorch_program, + ["et_vk.select_as_symint.default", "update_cache.default"], + ) + + def test_intermediate_cache_characterization(self) -> None: + exported, lowered = _export_intermediate_dynamic_update_cache() + self._assert_live_update_cache_symint(exported) + update_node = next( + node + for node in exported.graph_module.graph.nodes + if node.target == torch.ops.llama.update_cache.default + ) + self.assertEqual(torch.ops.aten.sigmoid.default, update_node.args[0].target) + self._assert_user_cache_writeback( + lowered.executorch_program, + [ + "aten.sigmoid.default", + "et_vk.select_as_symint.default", + "update_cache.default", + ], + ) + def export_update_cache_model(output_path: str) -> None: """Export an update_cache model to .pte for the native runtime test. @@ -192,5 +490,17 @@ def export_update_cache_negative(out_dir: str) -> None: print(f"Exported {name}.pte") -if __name__ == "__main__": - unittest.main() +def export_dynamic_update_cache(output_path: str) -> None: + """Export the one-PTE dynamic position/sequence/capacity fixture.""" + _, program = _export_dynamic_update_cache() + with open(output_path, "wb") as output: + output.write(program.buffer) + print(f"Exported {output_path}") + + +def export_intermediate_dynamic_update_cache(output_path: str) -> None: + """Export the post-fixpoint intermediate-value runtime fixture.""" + _, program = _export_intermediate_dynamic_update_cache() + with open(output_path, "wb") as output: + output.write(program.buffer) + print(f"Exported {output_path}") diff --git a/backends/webgpu/test/targets.bzl b/backends/webgpu/test/targets.bzl index 1f014dfa670..ecb4e086eeb 100644 --- a/backends/webgpu/test/targets.bzl +++ b/backends/webgpu/test/targets.bzl @@ -19,6 +19,23 @@ def define_common_targets(is_fbcode = False): ], ) + runtime.python_test( + name = "test_update_cache", + srcs = ["ops/test_update_cache.py"], + preload_deps = [ + "//executorch/extension/llm/custom_ops:custom_ops_aot_lib", + "//executorch/extension/llm/custom_ops:custom_ops_aot_py", + ], + deps = [ + "//caffe2:torch", + "//executorch/backends/vulkan/partitioner:vulkan_partitioner", + "//executorch/backends/vulkan/serialization:lib", + "//executorch/backends/vulkan:vulkan_preprocess", + "//executorch/examples/models/gemma4:webgpu_support", + "//executorch/exir:lib", + ], + ) + runtime.python_library( name = "tester", srcs = ["tester.py"], diff --git a/backends/webgpu/test/test_native_ci_contract.py b/backends/webgpu/test/test_native_ci_contract.py index 43b8ec6f56e..49e81227f55 100644 --- a/backends/webgpu/test/test_native_ci_contract.py +++ b/backends/webgpu/test/test_native_ci_contract.py @@ -73,3 +73,43 @@ def test_requires_dynamic_rope_fixture(self) -> None: self.assertIn("export_rope_hf_dynamic('${ROPE_HF_DIR}')", script) self.assertIn('WEBGPU_TEST_ROPE_HF_DIR="${ROPE_HF_DIR}"', script) self.assertIn('require_file "${ROPE_HF_DIR}/rope_hf_dynamic.pte"', script) + + def test_cat_2d_regressions_are_heavy_and_fail_closed(self) -> None: + backend = pathlib.Path(__file__).parents[1] + script = (backend / "scripts/test_webgpu_native_ci.sh").read_text() + dynamic_test = (backend / "test/native/test_dynamic_shape.cpp").read_text() + dispatch_test = (backend / "test/native/test_dispatch_2d.cpp").read_text() + cases = (backend / "test/op_tests/cases.py").read_text() + driver = (backend / "test/op_tests/op_test_driver.cpp").read_text() + + self.assertIn('require_file "${DYNAMIC_SHAPE_DIR}/dyn_cat_2d.pte"', script) + self.assertIn( + 'WEBGPU_TEST_HEAVY=1 $PYTHON_EXECUTABLE -c "\n' + "from executorch.backends.webgpu.test.ops.dynamic_shape." + "test_dynamic_shape_export import export_dynamic_shape_cases", + script, + ) + self.assertIn("DynamicShape.CatCrosses2dDispatchBoundary", script) + self.assertIn("[ PASSED ] 3 tests.", script) + self.assertIn('CAT_2D_TEST_DIR="/tmp/webgpu_cat_2d_test"', script) + self.assertIn("WEBGPU_TEST_HEAVY=1 $PYTHON_EXECUTABLE", script) + self.assertIn('--output "${CAT_2D_TEST_DIR}" --ops cat', script) + self.assertIn( + "run_with_required_device env WEBGPU_REQUIRE_DEVICE=1", script + ) + + self.assertIn("TEST(DynamicShape, CatCrosses2dDispatchBoundary)", dynamic_test) + self.assertIn( + "{kCat2dRows, kCat1dRows, kCat2dRows, kCat1dRows, kCat2dRows}", + dynamic_test, + ) + self.assertIn( + "TEST(CatDispatchGrid, RestoresBothDimensionsAcrossResize)", + dispatch_test, + ) + self.assertIn('name="folded_2d_full_output"', cases) + self.assertIn("inputs=((65536, 65), (65536, 1))", cases) + self.assertIn("heavy=True", cases) + self.assertIn('std::getenv("WEBGPU_REQUIRE_DEVICE")', driver) + self.assertIn("required_device_failure_exit_code", driver) + self.assertIn('std::printf("WebGPU device acquired (native)\\n")', driver) diff --git a/examples/models/gemma4/BUCK b/examples/models/gemma4/BUCK index 9b57f93bc56..f6b388bc04b 100644 --- a/examples/models/gemma4/BUCK +++ b/examples/models/gemma4/BUCK @@ -8,6 +8,48 @@ non_fbcode_target(_kind = define_common_targets,) fbcode_target(_kind = define_common_targets,) +fbcode_target(_kind = runtime.python_library, + name = "target_prefill_contract", + srcs = ["target_prefill_contract.py"], + _is_external_target = True, + base_module = "executorch.examples.models.gemma4", + resources = { + "generate_target_prefill_oracle.py": "generate_target_prefill_oracle.py", + }, + typing = True, + visibility = ["PUBLIC"], +) + +fbcode_target(_kind = runtime.python_library, + name = "target_prefill_producer", + srcs = ["generate_target_prefill_oracle.py"], + _is_external_target = True, + base_module = "executorch.examples.models.gemma4", + deps = [ + ":quant_utils", + ":target_prefill_contract", + ":text_decoder", + ":webgpu_support", + "//caffe2:torch", + ], + typing = True, + visibility = ["PUBLIC"], +) + +fbcode_target(_kind = runtime.python_binary, + name = "generate_target_prefill_oracle", + main_function = "executorch.examples.models.gemma4.generate_target_prefill_oracle.main", + preload_deps = [ + "//executorch/extension/llm/custom_ops:custom_ops_aot_lib", + "//executorch/extension/llm/custom_ops:custom_ops_aot_py", + "//executorch/kernels/quantized:aot_lib", + "//pytorch/ao/torchao/csrc/cpu/shared_kernels/embedding_xbit:op_embedding_xbit_aten", + "//pytorch/ao/torchao/csrc/cpu/shared_kernels/linear_8bit_act_xbit_weight:op_linear_8bit_act_xbit_weight_aten", + ], + deps = [":target_prefill_producer"], + typing = True, +) + define_webgpu_python_targets() # Text decoder module diff --git a/examples/models/gemma4/generate_target_prefill_oracle.py b/examples/models/gemma4/generate_target_prefill_oracle.py new file mode 100644 index 00000000000..c8a2066cdc8 --- /dev/null +++ b/examples/models/gemma4/generate_target_prefill_oracle.py @@ -0,0 +1,403 @@ +#!/usr/bin/env python3 +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Generate target-only eager evidence for Gemma 4 speculative decode.""" + +from __future__ import annotations + +import argparse +import gc +import hashlib +import json +import math +import socket +import sys + +from datetime import datetime, timezone +from pathlib import Path +from typing import Sequence + +import torch + +from executorch.examples.models.gemma4.target_prefill_contract import ( + canonical_json_bytes, + file_identity, + final_chunk_range, + prompt_plan_sha256, + prompt_tokens, + TARGET_PREFILL_ATOL, + TARGET_PREFILL_AUTHORITY, + TARGET_PREFILL_CHUNK_SIZE, + TARGET_PREFILL_CONTEXTS, + TARGET_PREFILL_ENVELOPE_KIND, + TARGET_PREFILL_RTOL, + TARGET_PREFILL_SCHEMA_VERSION, + validate_target_prefill_receipt, +) +from executorch.examples.models.gemma4.text_decoder.gemma4_attention import ( + Gemma4KVCache, +) + + +def _utc_now() -> str: + return datetime.now(timezone.utc).isoformat().replace("+00:00", "Z") + + +def chunk_ranges(context: int) -> list[tuple[int, int]]: + if context <= 0: + raise ValueError("target-prefill context must be positive") + ranges: list[tuple[int, int]] = [] + start = 0 + while start < context: + length = min(TARGET_PREFILL_CHUNK_SIZE, context - start) + ranges.append((start, length)) + start += length + return ranges + + +def _dtype_name(dtype: torch.dtype) -> str: + names = { + torch.float32: "float32", + torch.float16: "float16", + torch.bfloat16: "bfloat16", + torch.int64: "int64", + torch.int32: "int32", + } + if dtype not in names: + raise ValueError(f"unsupported target-prefill tensor dtype: {dtype}") + return names[dtype] + + +def _tensor_bytes(tensor: torch.Tensor) -> bytes: + value = tensor.detach().cpu().contiguous() + if value.dtype == torch.bfloat16: + value = value.view(torch.uint16) + return value.numpy().tobytes() + + +def _tensor_envelope(tensor: torch.Tensor) -> dict[str, object]: + if not torch.isfinite(tensor).all().item(): + raise ValueError("target-prefill tensor contains non-finite values") + encoded = _tensor_bytes(tensor) + return { + "byte_order": "little", + "dtype": _dtype_name(tensor.dtype), + "layout": "row_major_contiguous", + "sha256": hashlib.sha256(encoded).hexdigest(), + "shape": list(tensor.shape), + } + + +def _compare_av(fused: torch.Tensor, manual: torch.Tensor) -> dict[str, object]: + if ( + not torch.isfinite(fused).all().item() + or not torch.isfinite(manual).all().item() + ): + raise ValueError("target-prefill AV contains non-finite values") + torch.testing.assert_close( + fused, + manual, + rtol=TARGET_PREFILL_RTOL, + atol=TARGET_PREFILL_ATOL, + ) + difference = (fused.to(torch.float64) - manual.to(torch.float64)).abs() + max_abs = float(difference.max().item()) if difference.numel() else 0.0 + reference_rms = float( + torch.sqrt(torch.mean(fused.to(torch.float64).square())).item() + ) + error_rms = float(torch.sqrt(torch.mean(difference.square())).item()) + rel_rms = 0.0 if reference_rms == 0.0 else error_rms / reference_rms + if not math.isfinite(max_abs) or not math.isfinite(rel_rms): + raise ValueError("target-prefill AV metrics are non-finite") + return { + "atol": TARGET_PREFILL_ATOL, + "max_abs": max_abs, + "passed": True, + "rel_rms": rel_rms, + "rtol": TARGET_PREFILL_RTOL, + } + + +def _reset_target_kv_caches(model: torch.nn.Module) -> int: + count = 0 + with torch.no_grad(): + for module in model.modules(): + if not isinstance(module, Gemma4KVCache): + continue + module.k_cache.zero_() + module.v_cache.zero_() + count += 1 + if count == 0: + raise ValueError("target-prefill model has no Gemma4 KV caches") + return count + + +def _arm_config(use_custom_sdpa: bool) -> dict[str, object]: + return { + "dtype": "float32", + "enable_dynamic_shape": True, + "group_size": 128, + "max_seq_len": 8960, + "text_quantize": "8da4w+emb4", + "use_custom_sdpa": use_custom_sdpa, + "use_kv_cache": True, + "variant": "e2b", + } + + +def _load_target(checkpoint: Path, *, use_custom_sdpa: bool) -> torch.nn.Module: + from executorch.examples.models.gemma4.quant_utils import ( + apply_embedding_quantization, + apply_linear_quantization, + parse_quantize, + ) + from executorch.examples.models.gemma4.text_decoder.gemma4_config import ( + Gemma4Config, + ) + from executorch.examples.models.gemma4.text_decoder.gemma4_model import Gemma4Model + + config = Gemma4Config.from_config("e2b") + config.use_kv_cache = True + config.max_seq_len = 8960 + config.enable_dynamic_shape = True + config.use_custom_sdpa = use_custom_sdpa + model = Gemma4Model( + config=config, + checkpoint_path=str(checkpoint.resolve()), + dtype=torch.float32, + ).get_eager_model() + linear_quant, embedding_quant = parse_quantize("8da4w+emb4") + if embedding_quant: + model = apply_embedding_quantization(model, embedding_quant).eval() + if linear_quant: + model = apply_linear_quantization(model, linear_quant, group_size=128).eval() + return model.eval() + + +def _module_attribute(owner: object, name: str, label: str) -> torch.nn.Module: + value = getattr(owner, name, None) + if not isinstance(value, torch.nn.Module): + raise ValueError(f"target-prefill model has no {label}") + return value + + +def _resolve_capture_modules( + model: torch.nn.Module, +) -> tuple[torch.nn.Module, torch.nn.Module]: + text_model = _module_attribute(model, "model", "text model") + self_decoder = _module_attribute(text_model, "self_decoder", "self decoder") + layers = getattr(self_decoder, "layers", None) + if not isinstance(layers, torch.nn.ModuleList) or len(layers) == 0: + raise ValueError("target-prefill model has no decoder layers") + self_attn = _module_attribute(layers[0], "self_attn", "layer-0 attention") + o_proj = _module_attribute(self_attn, "o_proj", "layer-0 output projection") + lm_head = _module_attribute(text_model, "lm_head", "LM head") + return o_proj, lm_head + + +def _run_arm( + checkpoint: Path, + *, + use_custom_sdpa: bool, +) -> dict[int, dict[str, object]]: + model = _load_target(checkpoint, use_custom_sdpa=use_custom_sdpa) + layer0_o_proj, lm_head = _resolve_capture_modules(model) + captures: dict[str, torch.Tensor] = {} + + def capture_av(_module: torch.nn.Module, inputs: tuple[torch.Tensor, ...]) -> None: + if len(inputs) != 1: + raise ValueError("target-prefill o_proj hook expected one input") + captures["av"] = inputs[0].detach().clone() + + def capture_raw_logits( + _module: torch.nn.Module, + _inputs: tuple[torch.Tensor, ...], + output: torch.Tensor, + ) -> None: + captures["raw_logits"] = output.detach().clone() + + av_hook = layer0_o_proj.register_forward_pre_hook(capture_av) + logits_hook = lm_head.register_forward_hook(capture_raw_logits) + results: dict[int, dict[str, object]] = {} + try: + with torch.inference_mode(): + for context in TARGET_PREFILL_CONTEXTS: + reset_count = _reset_target_kv_caches(model) + captures.clear() + tokens = prompt_tokens(context) + post_logits: torch.Tensor | None = None + for start, length in chunk_ranges(context): + if length > TARGET_PREFILL_CHUNK_SIZE: + raise ValueError("target-prefill model call exceeds 512 tokens") + input_ids = torch.tensor( + [tokens[start : start + length]], dtype=torch.long + ) + input_pos = torch.arange(start, start + length, dtype=torch.long) + post_logits = model(input_ids, input_pos, None) + if post_logits is None: + raise ValueError("target-prefill context produced no logits") + raw_logits = captures.get("raw_logits") + av = captures.get("av") + if raw_logits is None or av is None: + raise ValueError( + "target-prefill hooks did not capture logits and AV" + ) + start, length = final_chunk_range(context) + if tuple(av.shape) != (1, length, 8 * 256): + raise ValueError("target-prefill layer-0 AV has the wrong shape") + results[context] = { + "av": av.reshape(1, length, 8, 256), + "cache_reset_count": reset_count, + "config": _arm_config(use_custom_sdpa), + "final_chunk_length": length, + "final_chunk_start": start, + "logits_post_softcap": post_logits.detach().clone(), + "logits_pre_softcap": raw_logits, + } + finally: + av_hook.remove() + logits_hook.remove() + del model + gc.collect() + return results + + +def _load_runtime_source_receipt( + path: Path, +) -> tuple[dict[str, object], dict[str, object], str]: + try: + document = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError) as error: + raise ValueError(f"invalid runtime-source receipt: {path}") from error + if not isinstance(document, dict): + raise ValueError("runtime-source receipt must be an object") + fbsource_commit = document.get("fbsource_commit") + if ( + not isinstance(fbsource_commit, str) + or len(fbsource_commit) != 40 + or any(character not in "0123456789abcdef" for character in fbsource_commit) + ): + raise ValueError("runtime-source receipt has an invalid fbsource commit") + return document, file_identity(path), fbsource_commit + + +def generate_target_prefill_receipt( + checkpoint_root: Path, + runtime_source_receipt: Path, + *, + command: Sequence[str], +) -> dict[str, object]: + from executorch.examples.models.gemma4.webgpu_artifact_manifest import ( + validate_export_identity, + ) + + started_at = _utc_now() + checkpoint_acquisition = dict(validate_export_identity(checkpoint_root)) + _, runtime_identity, fbsource_commit = _load_runtime_source_receipt( + runtime_source_receipt + ) + fused = _run_arm(checkpoint_root, use_custom_sdpa=True) + manual = _run_arm(checkpoint_root, use_custom_sdpa=False) + contexts: dict[str, object] = {} + for context in TARGET_PREFILL_CONTEXTS: + fused_result = fused[context] + manual_result = manual[context] + fused_av = fused_result["av"] + manual_av = manual_result["av"] + raw_logits = fused_result["logits_pre_softcap"] + post_logits = fused_result["logits_post_softcap"] + assert isinstance(fused_av, torch.Tensor) + assert isinstance(manual_av, torch.Tensor) + assert isinstance(raw_logits, torch.Tensor) + assert isinstance(post_logits, torch.Tensor) + raw_token = int(torch.argmax(raw_logits[:, -1, :], dim=-1).item()) + post_token = int(torch.argmax(post_logits[:, -1, :], dim=-1).item()) + if raw_token != post_token: + raise ValueError( + f"target-prefill softcap changes the token at context {context}" + ) + contexts[str(context)] = { + "arm_configs": { + "custom_sdpa_fused": fused_result["config"], + "manual_unfused": manual_result["config"], + }, + "cache_reset_counts": { + "custom_sdpa_fused": fused_result["cache_reset_count"], + "manual_unfused": manual_result["cache_reset_count"], + }, + "chunk_size": TARGET_PREFILL_CHUNK_SIZE, + "context": context, + "final_chunk_length": fused_result["final_chunk_length"], + "final_chunk_start": fused_result["final_chunk_start"], + "layer0_manual_unfused_vs_custom_sdpa_fused": { + "agreement": _compare_av(fused_av, manual_av), + "custom_sdpa_fused": _tensor_envelope(fused_av), + "manual_unfused": _tensor_envelope(manual_av), + }, + "logits_post_softcap": _tensor_envelope(post_logits), + "logits_pre_softcap": _tensor_envelope(raw_logits), + "prefill_token_post_softcap": post_token, + "prefill_token_raw": raw_token, + "prompt_plan_sha256": prompt_plan_sha256(context), + } + producer_path = Path(__file__) + producer_identity = file_identity(producer_path) + receipt: dict[str, object] = { + "authority": TARGET_PREFILL_AUTHORITY, + "checkpoint_acquisition": checkpoint_acquisition, + "contexts": contexts, + "envelope_kind": TARGET_PREFILL_ENVELOPE_KIND, + "producer": { + "fbsource_commit": fbsource_commit, + "runtime_source_receipt": runtime_identity, + "source_path": producer_path.name, + "source_sha256": producer_identity["sha256"], + }, + "run": { + "command": list(command), + "finished_at_utc": _utc_now(), + "host": socket.gethostname(), + "started_at_utc": started_at, + }, + "schema_version": TARGET_PREFILL_SCHEMA_VERSION, + } + validate_target_prefill_receipt( + receipt, + expected_checkpoint_acquisition=checkpoint_acquisition, + expected_producer_path=producer_path, + expected_producer_sha256=str(producer_identity["sha256"]), + expected_runtime_source_identity=runtime_identity, + expected_fbsource_commit=fbsource_commit, + ) + return receipt + + +def main(argv: Sequence[str] | None = None) -> int: + parser = argparse.ArgumentParser() + parser.add_argument("--checkpoints", type=Path, required=True) + parser.add_argument("--runtime-source-receipt", type=Path, required=True) + parser.add_argument("--contexts", required=True) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args(argv) + try: + contexts = tuple(int(value) for value in args.contexts.split(",")) + except ValueError as error: + raise ValueError("--contexts must be a comma-separated integer list") from error + if contexts != TARGET_PREFILL_CONTEXTS: + raise ValueError("target-prefill owner run requires the exact ten contexts") + command = ["generate_target_prefill_oracle", *(argv or sys.argv[1:])] + receipt = generate_target_prefill_receipt( + args.checkpoints, + args.runtime_source_receipt, + command=command, + ) + args.output.write_bytes(canonical_json_bytes(receipt)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/models/gemma4/target_prefill_contract.py b/examples/models/gemma4/target_prefill_contract.py new file mode 100644 index 00000000000..bd839c2bec9 --- /dev/null +++ b/examples/models/gemma4/target_prefill_contract.py @@ -0,0 +1,350 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Receipt contract for the Gemma 4 target-only prefill oracle.""" + +from __future__ import annotations + +import hashlib +import json +import math + +from pathlib import Path +from typing import Mapping, Sequence, TypeGuard + + +TARGET_PREFILL_SCHEMA_VERSION = 2 +TARGET_PREFILL_ENVELOPE_KIND = "target_prefill_v2" +TARGET_PREFILL_AUTHORITY = "target_only_eager" +TARGET_PREFILL_CHUNK_SIZE = 512 +TARGET_PREFILL_ATOL = 1e-4 +TARGET_PREFILL_RTOL = 1e-3 +TARGET_PREFILL_CONTEXTS = (128, 511, 512, 513, 514, 1024, 2048, 4096, 4097, 8192) +TARGET_PREFILL_VOCAB_SIZE = 262144 + +_TOP_LEVEL_KEYS = { + "authority", + "checkpoint_acquisition", + "contexts", + "envelope_kind", + "producer", + "run", + "schema_version", +} +_CONTEXT_KEYS = { + "arm_configs", + "cache_reset_counts", + "chunk_size", + "context", + "final_chunk_length", + "final_chunk_start", + "layer0_manual_unfused_vs_custom_sdpa_fused", + "logits_post_softcap", + "logits_pre_softcap", + "prefill_token_post_softcap", + "prefill_token_raw", + "prompt_plan_sha256", +} +_ARM_CONFIG_KEYS = { + "dtype", + "enable_dynamic_shape", + "group_size", + "max_seq_len", + "text_quantize", + "use_custom_sdpa", + "use_kv_cache", + "variant", +} +_TENSOR_KEYS = {"byte_order", "dtype", "layout", "sha256", "shape"} + + +def _is_int(value: object) -> TypeGuard[int]: + return isinstance(value, int) and not isinstance(value, bool) + + +def _is_number(value: object) -> bool: + return isinstance(value, (int, float)) and not isinstance(value, bool) + + +def _is_hex(value: object, length: int) -> bool: + return ( + isinstance(value, str) + and len(value) == length + and all(character in "0123456789abcdef" for character in value) + ) + + +def _mapping(value: object, label: str) -> Mapping[str, object]: + if not isinstance(value, dict): + raise ValueError(f"{label} must be an object") + return value + + +def _sequence(value: object, label: str) -> Sequence[object]: + if not isinstance(value, list): + raise ValueError(f"{label} must be a list") + return value + + +def _require_exact_keys( + value: Mapping[str, object], expected: set[str], label: str +) -> None: + if set(value) != expected: + raise ValueError( + f"{label} keys mismatch: expected {sorted(expected)}, got {sorted(value)}" + ) + + +def canonical_json_bytes(document: Mapping[str, object]) -> bytes: + return ( + json.dumps(document, indent=2, sort_keys=True, ensure_ascii=False) + "\n" + ).encode("utf-8") + + +def file_identity(path: Path) -> dict[str, object]: + if not path.is_file(): + raise ValueError(f"target-prefill input is not a regular file: {path}") + digest = hashlib.sha256() + byte_count = 0 + with path.open("rb") as source: + for block in iter(lambda: source.read(1024 * 1024), b""): + digest.update(block) + byte_count += len(block) + return {"bytes": byte_count, "sha256": digest.hexdigest()} + + +def prompt_tokens(context: int) -> list[int]: + if not _is_int(context) or context <= 0: + raise ValueError("target-prefill context must be a positive integer") + return [(index % (TARGET_PREFILL_VOCAB_SIZE - 1)) + 1 for index in range(context)] + + +def prompt_plan_sha256(context: int) -> str: + encoded = json.dumps(prompt_tokens(context), separators=(",", ":")).encode("utf-8") + return hashlib.sha256(encoded).hexdigest() + + +def final_chunk_range(context: int) -> tuple[int, int]: + if not _is_int(context) or context <= 0: + raise ValueError("target-prefill context must be a positive integer") + remainder = context % TARGET_PREFILL_CHUNK_SIZE + length = remainder if remainder else TARGET_PREFILL_CHUNK_SIZE + return context - length, length + + +def reviewed_producer_source_path() -> Path: + path = Path(__file__).with_name("generate_target_prefill_oracle.py") + if not path.is_file(): + raise ValueError("reviewed target-prefill producer resource is missing") + return path + + +def _validate_tensor_envelope( + value: object, + *, + label: str, + expected_shape: list[int], +) -> None: + envelope = _mapping(value, label) + _require_exact_keys(envelope, _TENSOR_KEYS, label) + if envelope.get("byte_order") != "little": + raise ValueError(f"{label} byte order must be little") + if envelope.get("dtype") != "float32": + raise ValueError(f"{label} dtype must be float32") + if envelope.get("layout") != "row_major_contiguous": + raise ValueError(f"{label} layout must be row_major_contiguous") + if not _is_hex(envelope.get("sha256"), 64): + raise ValueError(f"{label} has an invalid sha256") + shape = _sequence(envelope.get("shape"), f"{label} shape") + if list(shape) != expected_shape: + raise ValueError(f"{label} shape must be {expected_shape}") + + +def _validate_arm_configs(value: object) -> None: + configs = _mapping(value, "target-prefill arm configurations") + expected_names = {"custom_sdpa_fused", "manual_unfused"} + _require_exact_keys(configs, expected_names, "target-prefill arm configurations") + fused = _mapping(configs["custom_sdpa_fused"], "custom SDPA configuration") + manual = _mapping(configs["manual_unfused"], "manual SDPA configuration") + _require_exact_keys(fused, _ARM_CONFIG_KEYS, "custom SDPA configuration") + _require_exact_keys(manual, _ARM_CONFIG_KEYS, "manual SDPA configuration") + if fused.get("use_custom_sdpa") is not True: + raise ValueError("custom SDPA arm must enable use_custom_sdpa") + if manual.get("use_custom_sdpa") is not False: + raise ValueError("manual SDPA arm must disable use_custom_sdpa") + fused_without_route = dict(fused) + manual_without_route = dict(manual) + del fused_without_route["use_custom_sdpa"] + del manual_without_route["use_custom_sdpa"] + if fused_without_route != manual_without_route: + raise ValueError("target-prefill arm configurations differ beyond SDPA") + required = { + "dtype": "float32", + "enable_dynamic_shape": True, + "group_size": 128, + "max_seq_len": 8960, + "text_quantize": "8da4w+emb4", + "use_kv_cache": True, + "variant": "e2b", + } + if fused_without_route != required: + raise ValueError("target-prefill arm configuration is not production-shaped") + + +def _validate_context(value: object, expected_context: int) -> None: # noqa: C901 + context = _mapping(value, f"target-prefill context {expected_context}") + _require_exact_keys( + context, _CONTEXT_KEYS, f"target-prefill context {expected_context}" + ) + if context.get("context") != expected_context: + raise ValueError("target-prefill context value does not match its key") + if context.get("chunk_size") != TARGET_PREFILL_CHUNK_SIZE: + raise ValueError("target-prefill chunk size must be 512") + final_start, final_length = final_chunk_range(expected_context) + if ( + context.get("final_chunk_start") != final_start + or context.get("final_chunk_length") != final_length + ): + raise ValueError("target-prefill final chunk range mismatch") + if context.get("prompt_plan_sha256") != prompt_plan_sha256(expected_context): + raise ValueError("target-prefill prompt plan mismatch") + + _validate_arm_configs(context.get("arm_configs")) + reset_counts = _mapping( + context.get("cache_reset_counts"), "target-prefill cache reset counts" + ) + _require_exact_keys( + reset_counts, + {"custom_sdpa_fused", "manual_unfused"}, + "target-prefill cache reset counts", + ) + if any(not _is_int(count) or count <= 0 for count in reset_counts.values()): + raise ValueError("target-prefill cache reset counts must be positive") + + _validate_tensor_envelope( + context.get("logits_pre_softcap"), + label="pre-softcap logits", + expected_shape=[1, 1, TARGET_PREFILL_VOCAB_SIZE], + ) + _validate_tensor_envelope( + context.get("logits_post_softcap"), + label="post-softcap logits", + expected_shape=[1, 1, TARGET_PREFILL_VOCAB_SIZE], + ) + raw_token = context.get("prefill_token_raw") + post_token = context.get("prefill_token_post_softcap") + if ( + not _is_int(raw_token) + or raw_token < 0 + or raw_token >= TARGET_PREFILL_VOCAB_SIZE + ): + raise ValueError("target-prefill raw token is invalid") + if post_token != raw_token: + raise ValueError("target-prefill raw/post-softcap tokens differ") + + witness = _mapping( + context.get("layer0_manual_unfused_vs_custom_sdpa_fused"), + "target-prefill AV witness", + ) + _require_exact_keys( + witness, + {"agreement", "custom_sdpa_fused", "manual_unfused"}, + "target-prefill AV witness", + ) + av_shape = [1, final_length, 8, 256] + _validate_tensor_envelope( + witness.get("custom_sdpa_fused"), + label="custom SDPA AV tensor", + expected_shape=av_shape, + ) + _validate_tensor_envelope( + witness.get("manual_unfused"), + label="manual SDPA AV tensor", + expected_shape=av_shape, + ) + agreement = _mapping(witness.get("agreement"), "target-prefill AV agreement") + _require_exact_keys( + agreement, + {"atol", "max_abs", "passed", "rel_rms", "rtol"}, + "target-prefill AV agreement", + ) + if agreement.get("atol") != TARGET_PREFILL_ATOL: + raise ValueError("target-prefill AV agreement atol mismatch") + if agreement.get("rtol") != TARGET_PREFILL_RTOL: + raise ValueError("target-prefill AV agreement rtol mismatch") + if agreement.get("passed") is not True: + raise ValueError("target-prefill AV agreement must pass") + for metric in ("max_abs", "rel_rms"): + observed = agreement.get(metric) + if not _is_number(observed) or not math.isfinite(float(observed)): + raise ValueError(f"target-prefill AV agreement {metric} is not finite") + if float(observed) < 0: + raise ValueError(f"target-prefill AV agreement {metric} is negative") + + +def validate_target_prefill_receipt( # noqa: C901 + receipt: Mapping[str, object], + *, + expected_checkpoint_acquisition: Mapping[str, object], + expected_producer_path: Path, + expected_producer_sha256: str, + expected_runtime_source_identity: Mapping[str, object], + expected_fbsource_commit: str, +) -> None: + _require_exact_keys(receipt, _TOP_LEVEL_KEYS, "target-prefill receipt") + if receipt.get("schema_version") != TARGET_PREFILL_SCHEMA_VERSION: + raise ValueError("target-prefill schema version mismatch") + if receipt.get("envelope_kind") != TARGET_PREFILL_ENVELOPE_KIND: + raise ValueError("target-prefill envelope kind mismatch") + if receipt.get("authority") != TARGET_PREFILL_AUTHORITY: + raise ValueError("target-prefill authority mismatch") + if receipt.get("checkpoint_acquisition") != expected_checkpoint_acquisition: + raise ValueError("target-prefill checkpoint acquisition mismatch") + + producer = _mapping(receipt.get("producer"), "target-prefill producer") + _require_exact_keys( + producer, + { + "fbsource_commit", + "runtime_source_receipt", + "source_path", + "source_sha256", + }, + "target-prefill producer", + ) + if producer.get("source_path") != expected_producer_path.name: + raise ValueError("target-prefill producer source path mismatch") + if producer.get("source_sha256") != expected_producer_sha256: + raise ValueError("target-prefill producer source hash mismatch") + actual_identity = file_identity(expected_producer_path) + if actual_identity.get("sha256") != expected_producer_sha256: + raise ValueError("reviewed target-prefill producer bytes changed") + if producer.get("runtime_source_receipt") != expected_runtime_source_identity: + raise ValueError("target-prefill runtime source identity mismatch") + if producer.get("fbsource_commit") != expected_fbsource_commit: + raise ValueError("target-prefill fbsource commit mismatch") + if not _is_hex(expected_fbsource_commit, 40): + raise ValueError("expected target-prefill fbsource commit is invalid") + + run = _mapping(receipt.get("run"), "target-prefill run") + _require_exact_keys( + run, + {"command", "finished_at_utc", "host", "started_at_utc"}, + "target-prefill run", + ) + command = _sequence(run.get("command"), "target-prefill run command") + if not command or any(not isinstance(item, str) or not item for item in command): + raise ValueError("target-prefill run command must be nonempty strings") + for key in ("host", "started_at_utc", "finished_at_utc"): + if not isinstance(run.get(key), str) or not run[key]: + raise ValueError(f"target-prefill run {key} must be a nonempty string") + + contexts = _mapping(receipt.get("contexts"), "target-prefill contexts") + expected_keys = {str(context) for context in TARGET_PREFILL_CONTEXTS} + if set(contexts) != expected_keys: + raise ValueError("target-prefill receipt must contain the exact contexts") + for expected_context in TARGET_PREFILL_CONTEXTS: + _validate_context(contexts[str(expected_context)], expected_context) diff --git a/examples/models/gemma4/tests/targets.bzl b/examples/models/gemma4/tests/targets.bzl index c62ba6d2a64..fd1d12ecce4 100644 --- a/examples/models/gemma4/tests/targets.bzl +++ b/examples/models/gemma4/tests/targets.bzl @@ -14,6 +14,18 @@ def define_common_targets(is_fbcode = False): ], ) + fbcode_target(_kind = runtime.python_test, + name = "test_export_partitioners", + srcs = ["test_export_partitioners.py"], + deps = [ + "//caffe2:torch", + "//executorch/backends/vulkan:op_registry", + "//executorch/backends/vulkan/partitioner:vulkan_partitioner", + "//executorch/examples/models/gemma4:webgpu_support", + "//executorch/exir:lib", + ], + ) + fbcode_target(_kind = runtime.python_test, name = "test_webgpu_rewrite_pass", srcs = ["test_webgpu_rewrite_pass.py"], @@ -28,3 +40,55 @@ def define_common_targets(is_fbcode = False): "//executorch/exir:lib", ], ) + + fbcode_target(_kind = runtime.python_test, + name = "test_export_smoke", + srcs = ["test_export_smoke.py"], + deps = [ + "//caffe2:torch", + "//executorch/examples/models/gemma4:text_decoder", + ], + ) + + fbcode_target(_kind = runtime.python_test, + name = "test_gemma4_sdpa_host_contract", + srcs = ["test_gemma4_sdpa_host_contract.py"], + deps = [ + "//caffe2:torch", + "//executorch/backends/vulkan:custom_ops_lib", + "//executorch/examples/models/gemma4:webgpu_support", + ], + ) + + fbcode_target(_kind = runtime.python_test, + name = "test_selected_row_cross_decoder", + srcs = ["test_selected_row_cross_decoder.py"], + deps = [ + "//caffe2:torch", + "//executorch/examples/models/gemma4:text_decoder", + ], + ) + + fbcode_target(_kind = runtime.python_test, + name = "test_webgpu_artifact_manifest", + srcs = ["test_webgpu_artifact_manifest.py"], + typing = True, + deps = [ + "//executorch/examples/models/gemma4:webgpu_support", + ], + ) + + fbcode_target(_kind = runtime.python_test, + name = "test_gemma4_plain_wasm_contract", + srcs = ["test_gemma4_plain_wasm_contract.py"], + ) + + fbcode_target(_kind = runtime.python_test, + name = "test_generate_target_prefill_oracle", + srcs = ["test_generate_target_prefill_oracle.py"], + deps = [ + "//caffe2:torch", + "//executorch/examples/models/gemma4:target_prefill_contract", + "//executorch/examples/models/gemma4:target_prefill_producer", + ], + ) diff --git a/examples/models/gemma4/tests/test_export_partitioners.py b/examples/models/gemma4/tests/test_export_partitioners.py new file mode 100644 index 00000000000..8e769ca9179 --- /dev/null +++ b/examples/models/gemma4/tests/test_export_partitioners.py @@ -0,0 +1,53 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import unittest + +from executorch.backends.vulkan.op_registry import vulkan_supported_ops +from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.examples.models.gemma4.webgpu_partitioner import ( + _webgpu_allowlist, + build_webgpu_partitioner, +) +from executorch.exir.dialects._ops import ops as exir_ops + + +class ExportPartitionersTest(unittest.TestCase): + def test_plain_features_are_instance_scoped(self) -> None: + registry_before = dict(vulkan_supported_ops) + partitioner = build_webgpu_partitioner("8da4w+emb4") + + self.assertEqual(vulkan_supported_ops, registry_before) + self.assertIn( + exir_ops.edge.et_vk.apply_rotary_emb_hf_single.default, + partitioner._inner.extra_op_features, + ) + self.assertIn( + exir_ops.edge.et_vk.gemma4_sdpa.default, + partitioner._inner.extra_op_features, + ) + # Assert against the GLOBAL registry, not a default partitioner's + # instance map: that map is unconditionally empty, so the old form + # passed even if the op were globally registered. + self.assertNotIn( + exir_ops.edge.et_vk.apply_rotary_emb_hf_single.default, + vulkan_supported_ops, + ) + self.assertNotIn( + exir_ops.edge.et_vk.gemma4_sdpa.default, + vulkan_supported_ops, + ) + self.assertEqual(VulkanPartitioner().extra_op_features, {}) + + def test_restricted_allowlist_includes_symbolic_select(self) -> None: + allowlist = set(_webgpu_allowlist()) + self.assertIn(exir_ops.edge.et_vk.select_as_symint.default, allowlist) + self.assertNotIn(exir_ops.edge.aten.mm.default, allowlist) + self.assertNotIn(exir_ops.edge.aten.linear.default, allowlist) + + def test_emb8_fails_closed(self) -> None: + with self.assertRaisesRegex(ValueError, "emb8"): + build_webgpu_partitioner("8da4w+emb8") diff --git a/examples/models/gemma4/tests/test_export_smoke.py b/examples/models/gemma4/tests/test_export_smoke.py new file mode 100644 index 00000000000..576b5f18346 --- /dev/null +++ b/examples/models/gemma4/tests/test_export_smoke.py @@ -0,0 +1,39 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import unittest + +from types import SimpleNamespace +from unittest import mock + +from executorch.examples.models.gemma4.text_decoder.gemma4_model import Gemma4Model + + +def _model(max_seq_len: int = 8960) -> Gemma4Model: + model = Gemma4Model.__new__(Gemma4Model) + model.config = SimpleNamespace( + enable_dynamic_shape=True, + max_seq_len=max_seq_len, + use_kv_cache=True, + ) + return model + + +class ExportSmokeTest(unittest.TestCase): + def test_input_bound_is_independent_from_kv_capacity(self) -> None: + dim = object() + with mock.patch("torch.export.Dim", return_value=dim) as dim_factory: + dynamic_shapes = _model().get_dynamic_shapes(max_input_len=512) + + dim_factory.assert_called_once_with("seq_len", min=1, max=512) + self.assertIs(dynamic_shapes["input_ids"][1], dim) + self.assertIs(dynamic_shapes["input_pos"][0], dim) + + def test_input_bound_fails_closed(self) -> None: + for invalid in (1, 8960, 8961): + with self.subTest(max_input_len=invalid): + with self.assertRaisesRegex(ValueError, "max_input_len"): + _model().get_dynamic_shapes(max_input_len=invalid) diff --git a/examples/models/gemma4/tests/test_gemma4_plain_wasm_contract.py b/examples/models/gemma4/tests/test_gemma4_plain_wasm_contract.py new file mode 100644 index 00000000000..bb0f8515b30 --- /dev/null +++ b/examples/models/gemma4/tests/test_gemma4_plain_wasm_contract.py @@ -0,0 +1,183 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import re +import unittest + +from pathlib import Path + + +class Gemma4PlainWasmContractTest(unittest.TestCase): + @classmethod + def setUpClass(cls) -> None: + cls.model_root = Path(__file__).resolve().parents[1] + cls.executorch_root = cls.model_root.parents[2] + cls.runner_path = cls.model_root / "runner" / "gemma4_plain_wasm.cpp" + cls.runner = ( + cls.runner_path.read_text(encoding="utf-8") + if cls.runner_path.is_file() + else "" + ) + cls.backend_cmake = ( + cls.executorch_root / "backends" / "webgpu" / "CMakeLists.txt" + ).read_text(encoding="utf-8") + cls.model_cmake = (cls.model_root / "CMakeLists.txt").read_text( + encoding="utf-8" + ) + + def test_production_runner_exists(self) -> None: + self.assertTrue(self.runner_path.is_file()) + + def test_compare_abi_is_exported(self) -> None: + exports = set( + re.findall( + r"GEMMA4_WASM_EXPORT\s+(?:const char\*|void|int)\s+" + r"(et_[a-z0-9_]+)\s*\(", + self.runner, + ) + ) + self.assertTrue( + { + "et_init", + "et_load", + "et_unload", + "et_reset", + "et_prefill_batch", + "et_prefill_step", + "et_step", + "et_profile_enable", + "et_profile", + "et_get_last_prefill_token_count", + "et_get_route_contract_version", + "et_get_last_route_mask", + "et_get_last_route_conflict_count", + }.issubset(exports) + ) + self.assertNotIn("et_set_variant", exports) + + def test_compact_token_output_is_long_not_float_logits(self) -> None: + self.assertIn("output_tensor_meta(0)", self.runner) + self.assertGreaterEqual( + self.runner.count("ScalarType::Long"), + 4, + ) + self.assertIn("const_data_ptr()", self.runner) + self.assertNotIn("const_data_ptr()", self.runner) + self.assertIn("output.numel() != 1", self.runner) + + def test_load_requires_three_ordered_ptds(self) -> None: + self.assertIn("kExpectedPtdCount = 3", self.runner) + self.assertIn("ptd_paths.size() != kExpectedPtdCount", self.runner) + self.assertIn("load_webgpu_model", self.runner) + self.assertIn("std::move(ptd_paths)", self.runner) + + def test_reset_reloads_the_text_decoder_and_clears_observations(self) -> None: + reset_match = re.search( + r"GEMMA4_WASM_EXPORT\s+int\s+et_reset\(\)\s*\{(?P.*?)\n\}", + self.runner, + re.DOTALL, + ) + self.assertIsNotNone(reset_match) + body = reset_match.group("body") if reset_match is not None else "" + self.assertIn("unload_method(kMethodName)", body) + self.assertIn("load_method(kMethodName)", body) + self.assertIn("reset_runtime_observations()", body) + + observations_match = re.search( + r"void\s+reset_runtime_observations\(\)\s*\{(?P.*?)\n\}", + self.runner, + re.DOTALL, + ) + self.assertIsNotNone(observations_match) + observations = ( + observations_match.group("body") + if observations_match is not None + else "" + ) + self.assertIn("querypool->reset(0)", observations) + + def test_production_runner_has_no_variant_ab_switch(self) -> None: + self.assertNotIn("et_set_variant", self.runner) + self.assertNotIn("WEBGPU_VARIANT_", self.runner) + + def test_runner_has_no_dashboard_or_local_artifact_dependency(self) -> None: + forbidden = ( + "/home/", + "localhost", + "manifold", + "webgpu-e2e", + "webgpu_benchmark", + ) + for token in forbidden: + with self.subTest(token=token): + self.assertNotIn(token, self.runner.lower()) + + def test_cmake_and_buck_own_the_runner(self) -> None: + targets = (self.model_root / "targets.bzl").read_text(encoding="utf-8") + for build_file in (self.backend_cmake, targets): + with self.subTest(build_file=build_file): + self.assertIn("runner/gemma4_plain_wasm.cpp", build_file) + self.assertIn("webgpu_backend", build_file) + self.assertIn("webgpu_model_loader", build_file) + self.assertNotIn("gemma4_plain_wasm", self.model_cmake) + + def test_cmake_builds_a_compare_loadable_browser_module(self) -> None: + start = self.backend_cmake.index("add_executable(\n gemma4_plain_wasm") + end_marker = '"${CMAKE_CURRENT_BINARY_DIR}/browser_gemma4_plain"' + end = self.backend_cmake.index(end_marker, start) + len(end_marker) + cmake = self.backend_cmake[start:end] + required_link_contract = ( + "--use-port=emdawnwebgpu", + "-sASYNCIFY", + "-sALLOW_MEMORY_GROWTH=1", + "-sMAXIMUM_MEMORY=4GB", + "-sFORCE_FILESYSTEM=1", + "--no-entry", + "-sSTACK_SIZE=8388608", + "-sASYNCIFY_STACK_SIZE=1048576", + "-sMODULARIZE=1", + "-sEXPORT_NAME=createWebGPULlama", + 'OUTPUT_NAME "webgpu_llama"', + "browser_gemma4_plain", + ) + for option in required_link_contract: + with self.subTest(option=option): + self.assertIn(option, cmake) + self.assertNotIn("-sNO_ENTRY", cmake) + + runtime_methods = re.search( + r'-sEXPORTED_RUNTIME_METHODS=([^"\s]+)', cmake + ) + self.assertIsNotNone(runtime_methods) + self.assertEqual( + set(runtime_methods.group(1).split(",")) if runtime_methods else set(), + {"ccall", "cwrap", "FS", "HEAP32"}, + ) + expected_functions = { + "_et_init", + "_et_load", + "_et_unload", + "_et_reset", + "_et_step", + "_et_prefill_step", + "_et_prefill_batch", + "_et_profile_enable", + "_et_profile", + "_et_get_last_prefill_token_count", + "_et_get_route_contract_version", + "_et_get_last_route_mask", + "_et_get_last_route_conflict_count", + "_malloc", + "_free", + } + exported_functions = re.search(r'-sEXPORTED_FUNCTIONS=([^"\s]+)', cmake) + self.assertIsNotNone(exported_functions) + self.assertEqual( + set(exported_functions.group(1).split(",")) + if exported_functions + else set(), + expected_functions, + ) diff --git a/examples/models/gemma4/tests/test_gemma4_sdpa_host_contract.py b/examples/models/gemma4/tests/test_gemma4_sdpa_host_contract.py new file mode 100644 index 00000000000..f14d43fe6ab --- /dev/null +++ b/examples/models/gemma4/tests/test_gemma4_sdpa_host_contract.py @@ -0,0 +1,191 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""Host-side Gemma 4 `et_vk.gemma4_sdpa` ABI and architecture contract. + +This file pins the exported head geometry and the eager custom-op fence and +numerics. It does not execute WebGPU route selection. The handler's +occupancy-based QK route, masked-QK-elision predicate, `S_kv <= 4096` boundary, +and dynamic-resize flip belong to the Dawn native lane. +""" + +import copy +import unittest +from typing import Tuple + +import executorch.backends.vulkan.custom_ops_lib # noqa: F401 + +import torch +import torch.nn.functional as F +from executorch.examples.models.gemma4 import webgpu_artifact_manifest as wam + +GEMMA_HEADS = 8 +GEMMA_KV_HEADS = 1 +GEMMA_HEAD_DIMS = (256, 512) + +NEG_INF = float("-inf") + + +def _bshd( + s_q: int, s_kv: int, head_dim: int, *, batch: int = 1 +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Damped so `scale=1.0` logits stay off the softmax saturation floor.""" + generator = torch.Generator().manual_seed(0) + query = torch.randn(batch, s_q, GEMMA_HEADS, head_dim, generator=generator) * 0.05 + key = torch.randn(batch, s_kv, GEMMA_KV_HEADS, head_dim, generator=generator) * 0.05 + value = torch.randn(batch, s_kv, GEMMA_KV_HEADS, head_dim, generator=generator) + return query, key, value + + +def _reference( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + mask: torch.Tensor, +) -> torch.Tensor: + group = query.shape[2] // key.shape[2] + return F.scaled_dot_product_attention( + query.transpose(1, 2), + key.transpose(1, 2).repeat_interleave(group, dim=1), + value.transpose(1, 2).repeat_interleave(group, dim=1), + attn_mask=mask, + scale=1.0, + ).transpose(1, 2) + + +def _call( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + mask: torch.Tensor, + *, + start_pos: int = 0, + dropout_p: float = 0.0, + is_causal: bool = False, + scale: float = 1.0, +) -> torch.Tensor: + return torch.ops.et_vk.gemma4_sdpa.default( + query, key, value, start_pos, mask, dropout_p, is_causal, scale + ) + + +class Gemma4SdpaExportGeometryTest(unittest.TestCase): + def test_shipped_config_satisfies_the_architecture_validator(self) -> None: + # The validator only sees this config behind a full checkpoint root. + config = wam._load_json(wam._source_config_path()) + wam._validate_architecture(config, "source config") + + def test_architecture_validator_rejects_a_drifted_config(self) -> None: + config = copy.deepcopy(dict(wam._load_json(wam._source_config_path()))) + config["text_config"]["num_attention_heads"] = GEMMA_HEADS + 1 + with self.assertRaisesRegex(ValueError, "fingerprint mismatch"): + wam._validate_architecture(config, "source config") + + def test_numeric_cases_use_the_exported_head_geometry(self) -> None: + # The damped fixtures below use the exported architecture. + fingerprint = wam.ARCHITECTURE_FINGERPRINT + self.assertEqual(fingerprint["num_attention_heads"], GEMMA_HEADS) + self.assertEqual(fingerprint["num_key_value_heads"], GEMMA_KV_HEADS) + self.assertEqual( + (fingerprint["head_dim"], fingerprint["global_head_dim"]), + GEMMA_HEAD_DIMS, + ) + + +class Gemma4SdpaAbiTest(unittest.TestCase): + def test_every_fence_clause_fails_closed(self) -> None: + query, key, value = _bshd(1, 8, 256) + mask = torch.zeros(1, 8) + wrong_batch, _, _ = _bshd(1, 8, 256, batch=2) + wide_query, _, _ = _bshd(1, 8, 512) + grouped_key = torch.zeros(1, 8, 3, 256) + for name, expected, kwargs in [ + ("dropout", "dropout=0", {"dropout_p": 0.1}), + ("causal", "causal=false", {"is_causal": True}), + ("scale", "scale=1", {"scale": 0.5}), + ]: + with self.subTest(clause=name): + with self.assertRaisesRegex(ValueError, expected): + _call(query, key, value, mask, **kwargs) + + for name, expected, args in [ + ("rank-3 query", "BSHD", (query[0], key, value, mask)), + ("rank-3 key", "BSHD", (query, key[0], value, mask)), + ("rank-3 value", "BSHD", (query, key, value[0], mask)), + ( + "key/value mismatch", + "shapes do not match", + (query, key, value[:, :4], mask), + ), + ( + "batch mismatch", + "shapes do not match", + (wrong_batch, key, value, mask), + ), + ( + "head dim mismatch", + "grouped-query compatible", + (wide_query, key, value, mask), + ), + ( + "grouped-query mismatch", + "grouped-query compatible", + (query, grouped_key, grouped_key, mask), + ), + ( + "rank-4 mask", + r"rank-2 \[S_q, S_kv\] mask", + (query, key, value, mask.reshape(1, 1, 1, 8)), + ), + ( + "transposed mask", + r"rank-2 \[S_q, S_kv\] mask", + (query, key, value, mask.transpose(0, 1)), + ), + ]: + with self.subTest(clause=name): + with self.assertRaisesRegex(ValueError, expected): + _call(*args) + + +class Gemma4SdpaNumericsTest(unittest.TestCase): + def test_decode_geometry_matches_reference(self) -> None: + for head_dim in GEMMA_HEAD_DIMS: + with self.subTest(head_dim=head_dim): + query, key, value = _bshd(1, 32, head_dim) + mask = torch.zeros(1, 32) + torch.testing.assert_close( + _call(query, key, value, mask), + _reference(query, key, value, mask), + atol=1e-4, + rtol=1e-3, + ) + + def test_negative_infinity_mask_positions_are_inert(self) -> None: + # Every row keeps a live prefix: an all -inf row makes softmax NaN. + live = 5 + query, key, value = _bshd(1, 32, 256) + mask = torch.zeros(1, 32) + mask[:, live:] = NEG_INF + + got = _call(query, key, value, mask) + self.assertFalse(torch.isnan(got).any()) + torch.testing.assert_close( + got, + _reference(query, key[:, :live], value[:, :live], torch.zeros(1, live)), + atol=1e-4, + rtol=1e-3, + ) + + +class GenericSdpaFallbackTest(unittest.TestCase): + def test_generic_op_rejects_the_gemma4_positional_abi(self) -> None: + # `et_vk.sdpa` keeps its 5-value schema; numerics live in + # backends/webgpu/test/ops/test_et_vk_sdpa.py. + query, key, value = _bshd(1, 8, 256) + mask = torch.zeros(1, 8) + with self.assertRaises(RuntimeError): + torch.ops.et_vk.sdpa.default(query, key, value, 0, mask, 0.0, False, 1.0) diff --git a/examples/models/gemma4/tests/test_generate_target_prefill_oracle.py b/examples/models/gemma4/tests/test_generate_target_prefill_oracle.py new file mode 100644 index 00000000000..283ef4a8d79 --- /dev/null +++ b/examples/models/gemma4/tests/test_generate_target_prefill_oracle.py @@ -0,0 +1,618 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +from __future__ import annotations + +import ast +import hashlib +import importlib +import json +import math +import tempfile +import unittest + +from pathlib import Path +from typing import Any +from unittest import mock + +import torch + + +_PACKAGE = Path(__file__).parents[1] +_CONTRACT_MODULE = "executorch.examples.models.gemma4.target_prefill_contract" +_PRODUCER_MODULE = "executorch.examples.models.gemma4.generate_target_prefill_oracle" + + +def _modules() -> tuple[Any, Any]: + return importlib.import_module(_CONTRACT_MODULE), importlib.import_module( + _PRODUCER_MODULE + ) + + +def _tensor_envelope(shape: list[int], digest: str = "a" * 64) -> dict[str, object]: + return { + "byte_order": "little", + "dtype": "float32", + "layout": "row_major_contiguous", + "sha256": digest, + "shape": shape, + } + + +def _arm_config(use_custom_sdpa: bool) -> dict[str, object]: + return { + "dtype": "float32", + "enable_dynamic_shape": True, + "group_size": 128, + "max_seq_len": 8960, + "text_quantize": "8da4w+emb4", + "use_custom_sdpa": use_custom_sdpa, + "use_kv_cache": True, + "variant": "e2b", + } + + +class _OffsetProjection(torch.nn.Module): + def forward(self, value: torch.Tensor) -> torch.Tensor: + return value + 1000.0 + + +class _RecordingLmHead(torch.nn.Module): + def forward(self, value: torch.Tensor) -> torch.Tensor: + return value[..., :4] + + +class _RecordingTarget(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.state = 0 + self.calls: list[tuple[int, int, int, int, int]] = [] + self.model = torch.nn.Module() + self.model.self_decoder = torch.nn.Module() + layer = torch.nn.Module() + layer.self_attn = torch.nn.Module() + layer.self_attn.o_proj = _OffsetProjection() + self.model.self_decoder.layers = torch.nn.ModuleList([layer]) + self.model.lm_head = _RecordingLmHead() + + def forward( + self, + input_ids: torch.Tensor, + input_pos: torch.Tensor, + _mask: object, + ) -> torch.Tensor: + start = int(input_pos[0].item()) + length = int(input_ids.shape[1]) + self.calls.append( + ( + start, + length, + self.state, + int(input_ids[0, 0].item()), + int(input_ids[0, -1].item()), + ) + ) + av = torch.full( + (1, length, 8 * 256), + float(start + 1 + 100 * self.state), + dtype=torch.float32, + ) + projected = self.model.self_decoder.layers[0].self_attn.o_proj(av) + raw_logits = self.model.lm_head(projected) + self.state += 1 + return raw_logits + 50.0 + + +def _valid_receipt( + contract: Any, + checkpoint: dict[str, object], + runtime_identity: dict[str, object], + fbsource_commit: str, +) -> dict[str, object]: + producer_path = contract.reviewed_producer_source_path() + producer_digest = hashlib.sha256(producer_path.read_bytes()).hexdigest() + contexts: dict[str, object] = {} + for context in contract.TARGET_PREFILL_CONTEXTS: + start, length = contract.final_chunk_range(context) + contexts[str(context)] = { + "arm_configs": { + "custom_sdpa_fused": _arm_config(True), + "manual_unfused": _arm_config(False), + }, + "cache_reset_counts": { + "custom_sdpa_fused": 15, + "manual_unfused": 15, + }, + "chunk_size": 512, + "context": context, + "final_chunk_length": length, + "final_chunk_start": start, + "layer0_manual_unfused_vs_custom_sdpa_fused": { + "agreement": { + "atol": 1e-4, + "max_abs": 1e-5, + "passed": True, + "rel_rms": 1e-6, + "rtol": 1e-3, + }, + "custom_sdpa_fused": _tensor_envelope([1, length, 8, 256]), + "manual_unfused": _tensor_envelope([1, length, 8, 256], "b" * 64), + }, + "logits_post_softcap": _tensor_envelope([1, 1, 262144], "c" * 64), + "logits_pre_softcap": _tensor_envelope([1, 1, 262144], "d" * 64), + "prefill_token_post_softcap": 17, + "prefill_token_raw": 17, + "prompt_plan_sha256": contract.prompt_plan_sha256(context), + } + return { + "authority": "target_only_eager", + "checkpoint_acquisition": checkpoint, + "contexts": contexts, + "envelope_kind": "target_prefill_v2", + "producer": { + "fbsource_commit": fbsource_commit, + "runtime_source_receipt": runtime_identity, + "source_path": producer_path.name, + "source_sha256": producer_digest, + }, + "run": { + "command": ["generate_target_prefill_oracle", "--contexts", "all"], + "finished_at_utc": "2026-08-07T12:00:01Z", + "host": "test-host", + "started_at_utc": "2026-08-07T12:00:00Z", + }, + "schema_version": 2, + } + + +class TargetPrefillOracleConstructionTest(unittest.TestCase): + def test_public_contract_is_declared(self) -> None: + contract, producer = _modules() + contract_names = ( + "TARGET_PREFILL_CONTEXTS", + "canonical_json_bytes", + "file_identity", + "final_chunk_range", + "prompt_plan_sha256", + "reviewed_producer_source_path", + "validate_target_prefill_receipt", + ) + producer_names = ( + "_compare_av", + "_reset_target_kv_caches", + "_tensor_envelope", + "chunk_ranges", + ) + self.assertEqual( + [name for name in contract_names if not hasattr(contract, name)], [] + ) + self.assertEqual( + [name for name in producer_names if not hasattr(producer, name)], [] + ) + + def test_contract_and_producer_sources_exist(self) -> None: + self.assertTrue((_PACKAGE / "target_prefill_contract.py").is_file()) + self.assertTrue((_PACKAGE / "generate_target_prefill_oracle.py").is_file()) + + def test_producer_is_independent_of_speculative_export(self) -> None: + tree = ast.parse( + (_PACKAGE / "generate_target_prefill_oracle.py").read_text(encoding="utf-8") + ) + imports = [ + node.module + for node in ast.walk(tree) + if isinstance(node, ast.ImportFrom) and node.module is not None + ] + self.assertFalse(any("export_speculative" in module for module in imports)) + + def test_producer_has_no_pinned_sha256_literal(self) -> None: + tree = ast.parse( + (_PACKAGE / "generate_target_prefill_oracle.py").read_text(encoding="utf-8") + ) + pinned = [ + node.value + for node in ast.walk(tree) + if isinstance(node, ast.Constant) + and isinstance(node.value, str) + and len(node.value) == 64 + and all(character in "0123456789abcdef" for character in node.value) + ] + self.assertEqual(pinned, []) + + +class TargetPrefillContractTest(unittest.TestCase): + def setUp(self) -> None: + self.contract, _ = _modules() + self.checkpoint = { + "files": {"model.safetensors": {"bytes": 3, "sha256": "e" * 64}}, + "repo_id": "example/model", + "revision": "f" * 40, + } + self.runtime_identity = {"bytes": 17, "sha256": "1" * 64} + self.fbsource_commit = "2" * 40 + + def receipt(self) -> dict[str, object]: + return _valid_receipt( + self.contract, + self.checkpoint, + self.runtime_identity, + self.fbsource_commit, + ) + + def validate(self, receipt: dict[str, object]) -> None: + producer_path = self.contract.reviewed_producer_source_path() + self.contract.validate_target_prefill_receipt( + receipt, + expected_checkpoint_acquisition=self.checkpoint, + expected_producer_path=producer_path, + expected_producer_sha256=hashlib.sha256( + producer_path.read_bytes() + ).hexdigest(), + expected_runtime_source_identity=self.runtime_identity, + expected_fbsource_commit=self.fbsource_commit, + ) + + def test_context_set_and_boundary_chunk_ranges_are_exact(self) -> None: + self.assertEqual( + self.contract.TARGET_PREFILL_CONTEXTS, + (128, 511, 512, 513, 514, 1024, 2048, 4096, 4097, 8192), + ) + self.assertEqual(self.contract.final_chunk_range(511), (0, 511)) + self.assertEqual(self.contract.final_chunk_range(512), (0, 512)) + self.assertEqual(self.contract.final_chunk_range(513), (512, 1)) + self.assertEqual(self.contract.final_chunk_range(4097), (4096, 1)) + + def test_prompt_hash_uses_the_pinned_token_formula(self) -> None: + context = 514 + tokens = [(index % (262144 - 1)) + 1 for index in range(context)] + expected = hashlib.sha256( + json.dumps(tokens, separators=(",", ":")).encode("utf-8") + ).hexdigest() + self.assertEqual(self.contract.prompt_plan_sha256(context), expected) + + def test_canonical_document_has_sorted_utf8_and_trailing_newline(self) -> None: + self.assertEqual( + self.contract.canonical_json_bytes({"z": "é", "a": 1}), + '{\n "a": 1,\n "z": "é"\n}\n'.encode(), + ) + + def test_file_identity_hashes_the_actual_bytes(self) -> None: + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "receipt.json" + path.write_bytes(b"receipt\n") + self.assertEqual( + self.contract.file_identity(path), + { + "bytes": 8, + "sha256": hashlib.sha256(b"receipt\n").hexdigest(), + }, + ) + + def test_valid_receipt_passes(self) -> None: + self.validate(self.receipt()) + + def test_missing_boundary_context_fails(self) -> None: + receipt = self.receipt() + contexts = receipt["contexts"] + assert isinstance(contexts, dict) + del contexts["513"] + with self.assertRaisesRegex(ValueError, "exact contexts"): + self.validate(receipt) + + def test_wrong_producer_bytes_fail(self) -> None: + receipt = self.receipt() + producer = receipt["producer"] + assert isinstance(producer, dict) + producer["source_sha256"] = "3" * 64 + with self.assertRaisesRegex(ValueError, "producer source"): + self.validate(receipt) + + def test_runtime_source_identity_and_head_are_bound(self) -> None: + receipt = self.receipt() + producer = receipt["producer"] + assert isinstance(producer, dict) + producer["runtime_source_receipt"] = {"bytes": 17, "sha256": "4" * 64} + with self.assertRaisesRegex(ValueError, "runtime source"): + self.validate(receipt) + receipt = self.receipt() + producer = receipt["producer"] + assert isinstance(producer, dict) + producer["fbsource_commit"] = "5" * 40 + with self.assertRaisesRegex(ValueError, "fbsource commit"): + self.validate(receipt) + + def test_av_tolerance_pass_shape_and_finite_metrics_are_enforced(self) -> None: + mutations: list[tuple[str, object]] = [ + ("atol", 1e-3), + ("rtol", 1e-2), + ("passed", False), + ("max_abs", math.inf), + ("rel_rms", math.nan), + ] + for key, value in mutations: + with self.subTest(key=key): + receipt = self.receipt() + context = receipt["contexts"]["513"] # type: ignore[index] + witness = context[ # type: ignore[index] + "layer0_manual_unfused_vs_custom_sdpa_fused" + ] + witness["agreement"][key] = value # type: ignore[index] + with self.assertRaises(ValueError): + self.validate(receipt) + receipt = self.receipt() + context = receipt["contexts"]["513"] # type: ignore[index] + witness = context[ # type: ignore[index] + "layer0_manual_unfused_vs_custom_sdpa_fused" + ] + witness["custom_sdpa_fused"]["shape"] = [1, 2, 8, 256] # type: ignore[index] + with self.assertRaisesRegex(ValueError, "AV tensor"): + self.validate(receipt) + + def test_arm_configs_may_differ_only_in_custom_sdpa(self) -> None: + receipt = self.receipt() + context = receipt["contexts"]["128"] # type: ignore[index] + context["arm_configs"]["manual_unfused"]["max_seq_len"] = 1024 # type: ignore[index] + with self.assertRaisesRegex(ValueError, "arm configurations"): + self.validate(receipt) + + def test_raw_and_post_softcap_tokens_must_match(self) -> None: + receipt = self.receipt() + context = receipt["contexts"]["128"] # type: ignore[index] + context["prefill_token_post_softcap"] = 18 # type: ignore[index] + with self.assertRaisesRegex(ValueError, "softcap"): + self.validate(receipt) + + def test_boolean_counts_and_tokens_are_not_integers(self) -> None: + receipt = self.receipt() + context = receipt["contexts"]["128"] # type: ignore[index] + context["cache_reset_counts"]["custom_sdpa_fused"] = True # type: ignore[index] + with self.assertRaisesRegex(ValueError, "reset counts"): + self.validate(receipt) + + receipt = self.receipt() + context = receipt["contexts"]["128"] # type: ignore[index] + context["prefill_token_raw"] = True # type: ignore[index] + context["prefill_token_post_softcap"] = True # type: ignore[index] + with self.assertRaisesRegex(ValueError, "raw token"): + self.validate(receipt) + + def test_unqualified_logits_digest_is_rejected(self) -> None: + receipt = self.receipt() + context = receipt["contexts"]["128"] # type: ignore[index] + context["logits_sha256"] = "6" * 64 # type: ignore[index] + with self.assertRaisesRegex(ValueError, "keys mismatch"): + self.validate(receipt) + + +class TargetPrefillProducerHelpersTest(unittest.TestCase): + def setUp(self) -> None: + _, self.producer = _modules() + + def _run_recording_arm( + self, contexts: tuple[int, ...] + ) -> tuple[ + _RecordingTarget, + list[int], + dict[int, dict[str, object]], + ]: + model = _RecordingTarget() + reset_states: list[int] = [] + + def reset(target: torch.nn.Module) -> int: + if target is not model: + raise AssertionError("reset received the wrong target") + reset_states.append(model.state) + model.state = 0 + return 7 + + with mock.patch.object( + self.producer, "_load_target", return_value=model + ), mock.patch.object( + self.producer, "_reset_target_kv_caches", side_effect=reset + ), mock.patch.object( + self.producer, "TARGET_PREFILL_CONTEXTS", contexts + ): + results = self.producer._run_arm( + Path("model.safetensors"), use_custom_sdpa=True + ) + return model, reset_states, results + + def test_run_arm_chunks_calls_and_resets_state_between_contexts(self) -> None: + forward_model, forward_resets, forward = self._run_recording_arm((511, 513)) + reverse_model, reverse_resets, reverse = self._run_recording_arm((513, 511)) + + self.assertEqual(forward_resets, [0, 1]) + self.assertEqual(reverse_resets, [0, 2]) + self.assertEqual( + forward_model.calls, + [ + (0, 511, 0, 1, 511), + (0, 512, 0, 1, 512), + (512, 1, 1, 513, 513), + ], + ) + self.assertEqual( + reverse_model.calls, + [ + (0, 512, 0, 1, 512), + (512, 1, 1, 513, 513), + (0, 511, 0, 1, 511), + ], + ) + for context in (511, 513): + with self.subTest(context=context): + self.assertEqual(forward[context]["cache_reset_count"], 7) + self.assertEqual(reverse[context]["cache_reset_count"], 7) + self.assertEqual( + forward[context]["final_chunk_start"], + self.producer.final_chunk_range(context)[0], + ) + self.assertEqual( + forward[context]["final_chunk_length"], + self.producer.final_chunk_range(context)[1], + ) + for key in ("av", "logits_pre_softcap", "logits_post_softcap"): + self.assertTrue( + torch.equal(forward[context][key], reverse[context][key]) + ) + + def test_generate_receipt_passes_checkpoint_directory_to_weight_loader( + self, + ) -> None: + contract, _ = _modules() + manifest = importlib.import_module( + "executorch.examples.models.gemma4.webgpu_artifact_manifest" + ) + convert_weights = importlib.import_module( + "executorch.examples.models.gemma4.text_decoder.convert_weights" + ) + config_module = importlib.import_module( + "executorch.examples.models.gemma4.text_decoder.gemma4_config" + ) + + with tempfile.TemporaryDirectory() as directory: + checkpoint_root = Path(directory) / "checkpoint" + checkpoint_root.mkdir() + (checkpoint_root / "model.safetensors").write_bytes(b"fixture") + runtime_receipt = Path(directory) / "runtime-source.json" + runtime_receipt.write_text( + json.dumps({"fbsource_commit": "2" * 40}), encoding="utf-8" + ) + + logits = torch.zeros( + (1, 1, contract.TARGET_PREFILL_VOCAB_SIZE), + dtype=torch.float32, + ) + + def run_loader_boundary( + checkpoint: Path, *, use_custom_sdpa: bool + ) -> dict[int, dict[str, object]]: + config = config_module.Gemma4Config.from_config("e2b") + convert_weights.convert_hf_to_custom(str(checkpoint), config) + return { + context: { + "av": torch.zeros( + (1, self.producer.final_chunk_range(context)[1], 8, 256), + dtype=torch.float32, + ), + "cache_reset_count": 1, + "config": _arm_config(use_custom_sdpa), + "final_chunk_length": self.producer.final_chunk_range(context)[ + 1 + ], + "final_chunk_start": self.producer.final_chunk_range(context)[ + 0 + ], + "logits_post_softcap": logits, + "logits_pre_softcap": logits, + } + for context in self.producer.TARGET_PREFILL_CONTEXTS + } + + safe_open = mock.MagicMock() + safe_open.return_value.__enter__.return_value.keys.return_value = () + with mock.patch.object( + manifest, + "validate_export_identity", + return_value=manifest.CHECKPOINT_ACQUISITION, + ), mock.patch.object( + self.producer, "_run_arm", side_effect=run_loader_boundary + ), mock.patch( + "safetensors.safe_open", safe_open + ): + receipt = self.producer.generate_target_prefill_receipt( + checkpoint_root, + runtime_receipt, + command=["generate_target_prefill_oracle"], + ) + + self.assertEqual( + receipt["checkpoint_acquisition"], manifest.CHECKPOINT_ACQUISITION + ) + + def test_run_arm_uses_preprojection_av_and_presoftcap_logits(self) -> None: + _, _, results = self._run_recording_arm((513,)) + result = results[513] + av = result["av"] + raw_logits = result["logits_pre_softcap"] + post_logits = result["logits_post_softcap"] + + self.assertIsInstance(av, torch.Tensor) + self.assertIsInstance(raw_logits, torch.Tensor) + self.assertIsInstance(post_logits, torch.Tensor) + self.assertTrue(torch.equal(av, torch.full((1, 1, 8, 256), 613.0))) + self.assertTrue(torch.equal(raw_logits, torch.full((1, 1, 4), 1613.0))) + self.assertTrue(torch.equal(post_logits, torch.full((1, 1, 4), 1663.0))) + + def test_chunk_ranges_never_exceed_512(self) -> None: + for context in (511, 512, 513, 4096, 4097, 8192): + ranges = self.producer.chunk_ranges(context) + self.assertEqual(ranges[-1], self.producer.final_chunk_range(context)) + self.assertTrue(all(length <= 512 for _, length in ranges)) + self.assertEqual(sum(length for _, length in ranges), context) + + def test_tensor_envelope_canonicalizes_noncontiguous_tensors(self) -> None: + tensor = torch.arange(12, dtype=torch.float32).view(3, 4).t() + envelope = self.producer._tensor_envelope(tensor) + expected_bytes = tensor.detach().cpu().contiguous().numpy().tobytes() + self.assertEqual(envelope["shape"], [4, 3]) + self.assertEqual(envelope["dtype"], "float32") + self.assertEqual(envelope["sha256"], hashlib.sha256(expected_bytes).hexdigest()) + + def test_av_comparison_uses_tolerance_not_hash_equality(self) -> None: + fused = torch.ones((1, 2, 8, 256), dtype=torch.float32) + close = fused + 1e-5 + agreement = self.producer._compare_av(fused, close) + self.assertTrue(agreement["passed"]) + self.assertNotEqual( + self.producer._tensor_envelope(fused)["sha256"], + self.producer._tensor_envelope(close)["sha256"], + ) + with self.assertRaises(AssertionError): + self.producer._compare_av(fused, fused + 0.1) + + def test_reset_zeros_every_real_kv_cache(self) -> None: + from executorch.examples.models.gemma4.text_decoder.gemma4_attention import ( + Gemma4KVCache, + ) + + def cache() -> Any: + value = Gemma4KVCache.__new__(Gemma4KVCache) + torch.nn.Module.__init__(value) + value.register_buffer("k_cache", torch.ones((1, 2, 1, 2))) + value.register_buffer("v_cache", torch.ones((1, 2, 1, 2)) * 2) + return value + + model = torch.nn.Sequential(cache(), cache()) + self.assertEqual(self.producer._reset_target_kv_caches(model), 2) + for module in model.modules(): + if isinstance(module, Gemma4KVCache): + self.assertEqual(torch.count_nonzero(module.k_cache).item(), 0) + self.assertEqual(torch.count_nonzero(module.v_cache).item(), 0) + + def test_reset_fails_when_model_has_no_kv_cache(self) -> None: + with self.assertRaisesRegex(ValueError, "no Gemma4 KV caches"): + self.producer._reset_target_kv_caches(torch.nn.Linear(2, 2)) + + def test_capture_modules_are_runtime_validated(self) -> None: + with self.assertRaisesRegex(ValueError, "text model"): + self.producer._resolve_capture_modules(torch.nn.Module()) + + model = torch.nn.Module() + model.model = torch.nn.Module() + model.model.self_decoder = torch.nn.Module() + model.model.self_decoder.layers = torch.nn.ModuleList([torch.nn.Module()]) + layer = model.model.self_decoder.layers[0] + layer.self_attn = torch.nn.Module() + layer.self_attn.o_proj = torch.nn.Linear(2, 2) + model.model.lm_head = torch.nn.Linear(2, 2) + o_proj, lm_head = self.producer._resolve_capture_modules(model) + self.assertIs(o_proj, layer.self_attn.o_proj) + self.assertIs(lm_head, model.model.lm_head) + + +if __name__ == "__main__": + unittest.main() diff --git a/examples/models/gemma4/tests/test_selected_row_cross_decoder.py b/examples/models/gemma4/tests/test_selected_row_cross_decoder.py new file mode 100644 index 00000000000..3a2f125b79e --- /dev/null +++ b/examples/models/gemma4/tests/test_selected_row_cross_decoder.py @@ -0,0 +1,78 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import unittest + +import torch + +from executorch.examples.models.gemma4.text_decoder.gemma4_transformer import ( + Gemma4TextModel, +) + + +class _SelfDecoder(torch.nn.Module): + def forward(self, input_ids, input_pos=None, inputs_embeds=None): + del input_pos, inputs_embeds + seq_len = input_ids.shape[1] + hidden = torch.arange(seq_len * 3, dtype=torch.float32).reshape(1, seq_len, 3) + per_layer = hidden.unsqueeze(0).repeat(2, 1, 1, 1) + return hidden, per_layer, {} + + +class _CrossDecoder(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.last_hidden_shape = None + self.last_per_layer_shape = None + self.last_query_start_pos = None + + def forward( + self, + hidden_states, + per_layer_inputs, + shared_kv, + input_pos=None, + query_start_pos=None, + ): + del shared_kv, input_pos + self.last_hidden_shape = tuple(hidden_states.shape) + self.last_per_layer_shape = tuple(per_layer_inputs.shape) + self.last_query_start_pos = query_start_pos + return hidden_states + + +def _model() -> Gemma4TextModel: + model = Gemma4TextModel.__new__(Gemma4TextModel) + torch.nn.Module.__init__(model) + model.self_decoder = _SelfDecoder() + model.cross_decoder = _CrossDecoder() + model.norm = torch.nn.Identity() + model.lm_head = torch.nn.Identity() + model.final_logit_softcapping = 0.0 + return model + + +class SelectedRowCrossDecoderTest(unittest.TestCase): + def test_generation_narrows_before_cross_decoder(self) -> None: + model = _model() + input_ids = torch.ones((1, 4), dtype=torch.long) + logits = model(input_ids, input_pos=torch.arange(4)) + + self.assertEqual(tuple(logits.shape), (1, 1, 3)) + self.assertEqual(model.cross_decoder.last_hidden_shape, (1, 1, 3)) + self.assertEqual(model.cross_decoder.last_per_layer_shape, (2, 1, 1, 3)) + self.assertEqual(model.cross_decoder.last_query_start_pos, 3) + torch.testing.assert_close(logits, torch.tensor([[[9.0, 10.0, 11.0]]])) + + def test_non_generation_keeps_full_cross_decoder_input(self) -> None: + model = _model() + input_ids = torch.ones((1, 4), dtype=torch.long) + logits = model(input_ids) + + self.assertEqual(tuple(logits.shape), (1, 1, 3)) + self.assertEqual(model.cross_decoder.last_hidden_shape, (1, 4, 3)) + self.assertEqual(model.cross_decoder.last_per_layer_shape, (2, 1, 4, 3)) + self.assertIsNone(model.cross_decoder.last_query_start_pos) diff --git a/examples/models/gemma4/tests/test_webgpu_artifact_manifest.py b/examples/models/gemma4/tests/test_webgpu_artifact_manifest.py new file mode 100644 index 00000000000..68d05d9e604 --- /dev/null +++ b/examples/models/gemma4/tests/test_webgpu_artifact_manifest.py @@ -0,0 +1,405 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import copy +import hashlib +import json +import tempfile +import unittest + +from pathlib import Path +from typing import Any +from unittest import mock + +from executorch.examples.models.gemma4 import ( + webgpu_artifact_manifest as gemma4_manifest, +) + +from executorch.examples.models.gemma4.webgpu_artifact_manifest import ( + ARCHITECTURE_FINGERPRINT, + create_plain_manifest, + validate_plain_manifest, +) + + +def _set_digest(value: object) -> str: + return hashlib.sha256( + json.dumps(value, separators=(",", ":"), sort_keys=True).encode("utf-8") + ).hexdigest() + + +def _test_source_manifest() -> dict[str, Any]: + logical_path = "examples/models/gemma4/webgpu_artifact_manifest.py" + identity = {"bytes": 7, "sha256": "3" * 64} + return { + "checkouts": { + "fbsource": {"clean": True, "head": "1" * 40}, + "oss": {"clean": True, "head": "2" * 40}, + }, + "file_set_sha256": _set_digest([logical_path]), + "files": [ + { + "copies": { + "fbcode": { + **identity, + "path": f"fbcode/executorch/{logical_path}", + }, + "oss": {**identity, "path": logical_path}, + "xplat": { + **identity, + "path": f"xplat/executorch/{logical_path}", + }, + }, + "path": logical_path, + } + ], + "schema_version": 1, + } + + +def _test_wgsl_manifest() -> dict[str, Any]: + roles = ( + ("runtime/WebGPUShaderRegistry.cpp", "global_registry"), + ("runtime/ops/add/binary_add.wgsl", "wgsl"), + ("runtime/ops/add/binary_add_wgsl.h", "generated_header"), + ("scripts/gen_wgsl_headers.py", "generator"), + ) + files = [ + {"bytes": 7, "path": path, "role": role, "sha256": "4" * 64} + for path, role in roles + ] + return { + "fbsource_commit": "1" * 40, + "file_set_sha256": _set_digest( + [{"path": path, "role": role} for path, role in roles] + ), + "files": files, + "orphans": [], + "schema_version": 1, + } + + +def _sealed_source_receipt() -> dict[str, Any]: + return { + "fbsource_commit": "1" * 40, + "oss_commit": "2" * 40, + "schema_version": 3, + "source_current": True, + "source_manifest": _test_source_manifest(), + "verification": { + "source_checkout": "verified", + "wgsl_codegen": "verified", + }, + "wgsl_manifest": _test_wgsl_manifest(), + } + + +class WebGPUArtifactManifestTest(unittest.TestCase): + def setUp(self) -> None: + self._temporary_directory = tempfile.TemporaryDirectory() + self.addCleanup(self._temporary_directory.cleanup) + self.root = Path(self._temporary_directory.name) + self.pte = self.root / "model.pte" + self.ptds = [self.root / f"constants_{index}.ptd" for index in range(3)] + self.source_receipt = self.root / "source_receipt.json" + self.pte.write_bytes(b"shared") + self.ptds[0].write_bytes(b"shared") + self.ptds[1].write_bytes(b"ptd-one") + self.ptds[2].write_bytes(b"ptd-two") + self.source_receipt.write_text( + json.dumps(_sealed_source_receipt()), + encoding="utf-8", + ) + self.manifest = create_plain_manifest( + self.root, + { + "pte": Path(self.pte.name), + "source": Path(self.source_receipt.name), + }, + [Path(path.name) for path in self.ptds], + ) + + def test_round_trip_and_order(self) -> None: + validate_plain_manifest(self.root, self.manifest) + self.assertEqual( + self.manifest["ptd_order"], [path.name for path in self.ptds] + ) + self.assertEqual( + self.manifest["model"]["architecture"], ARCHITECTURE_FINGERPRINT + ) + + def test_rejects_wrong_bytes_and_hash(self) -> None: + self.ptds[1].write_bytes(b"changed") + with self.assertRaisesRegex(ValueError, "byte count|SHA-256"): + validate_plain_manifest(self.root, self.manifest) + + def test_rejects_missing_and_extra_artifacts(self) -> None: + self.ptds[2].unlink() + with self.assertRaises((FileNotFoundError, ValueError)): + validate_plain_manifest(self.root, self.manifest) + self.ptds[2].write_bytes(b"ptd-two") + (self.root / "extra.bin").write_bytes(b"extra") + with self.assertRaisesRegex(ValueError, "missing or extra"): + validate_plain_manifest(self.root, self.manifest) + + def test_rejects_internal_symlink(self) -> None: + self.pte.unlink() + self.pte.symlink_to(self.ptds[0].name) + with self.assertRaisesRegex(ValueError, "symlink"): + validate_plain_manifest(self.root, self.manifest) + + def test_rejects_architecture_mutation(self) -> None: + mutated = copy.deepcopy(self.manifest) + mutated["model"]["architecture"]["hidden_size"] += 1 + with self.assertRaisesRegex(ValueError, "architecture"): + validate_plain_manifest(self.root, mutated) + + def test_rejects_unsealed_source_receipt(self) -> None: + receipt = json.loads(self.source_receipt.read_text(encoding="utf-8")) + receipt["source_current"] = False + self.source_receipt.write_text(json.dumps(receipt), encoding="utf-8") + with self.assertRaisesRegex(ValueError, "source-current"): + create_plain_manifest( + self.root, + { + "pte": Path(self.pte.name), + "source": Path(self.source_receipt.name), + }, + [Path(path.name) for path in self.ptds], + ) + + def test_rejects_source_mirror_mismatch(self) -> None: + source_manifest = _test_source_manifest() + source_manifest["files"][0]["copies"]["oss"]["sha256"] = "9" * 64 + with self.assertRaisesRegex(ValueError, "mirror/OSS identity mismatch"): + gemma4_manifest.validate_source_manifest(source_manifest) + + def test_rejects_source_file_set_digest_mismatch(self) -> None: + source_manifest = _test_source_manifest() + source_manifest["file_set_sha256"] = "9" * 64 + with self.assertRaisesRegex(ValueError, "file-set identity mismatch"): + gemma4_manifest.validate_source_manifest(source_manifest) + + def test_rejects_invalid_checkout_head(self) -> None: + source_manifest = _test_source_manifest() + source_manifest["checkouts"]["fbsource"]["head"] = "short" + with self.assertRaisesRegex(ValueError, "checkout identity is invalid"): + gemma4_manifest.validate_source_manifest(source_manifest) + + def test_rejects_declared_wgsl_orphan(self) -> None: + wgsl_manifest = _test_wgsl_manifest() + wgsl_manifest["orphans"] = ["runtime/ops/add/orphan_wgsl.h"] + with self.assertRaisesRegex(ValueError, "orphan"): + gemma4_manifest.validate_wgsl_manifest(wgsl_manifest) + + def test_rejects_incomplete_source_verification(self) -> None: + receipt = _sealed_source_receipt() + receipt["verification"] = {"source_checkout": "verified"} + self.source_receipt.write_text(json.dumps(receipt), encoding="utf-8") + with self.assertRaisesRegex(ValueError, "verification is incomplete"): + create_plain_manifest( + self.root, + { + "pte": Path(self.pte.name), + "source": Path(self.source_receipt.name), + }, + [Path(path.name) for path in self.ptds], + ) + + def test_rejects_receipt_wgsl_head_mismatch(self) -> None: + receipt = _sealed_source_receipt() + receipt["wgsl_manifest"]["fbsource_commit"] = "9" * 40 + self.source_receipt.write_text(json.dumps(receipt), encoding="utf-8") + with self.assertRaisesRegex(ValueError, "mismatched WGSL checkout"): + create_plain_manifest( + self.root, + { + "pte": Path(self.pte.name), + "source": Path(self.source_receipt.name), + }, + [Path(path.name) for path in self.ptds], + ) + + +class SourceClosureManifestTest(unittest.TestCase): + def setUp(self) -> None: + self._temporary_directory = tempfile.TemporaryDirectory() + self.addCleanup(self._temporary_directory.cleanup) + self.root = Path(self._temporary_directory.name) + self.fbsource_root = self.root / "fbsource" + self.oss_root = self.root / "oss" + self.fbsource_root.mkdir() + self.oss_root.mkdir() + self.logical_path = "examples/models/gemma4/webgpu_artifact_manifest.py" + for path in ( + self.fbsource_root / "fbcode/executorch" / self.logical_path, + self.fbsource_root / "xplat/executorch" / self.logical_path, + self.oss_root / self.logical_path, + ): + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(b"source") + + def _create_source_manifest(self) -> dict[str, Any]: + def snapshot(_root: Path, kind: str) -> dict[str, object]: + return { + "clean": True, + "head": ("1" if kind == "fbsource" else "2") * 40, + } + + with mock.patch.object( + gemma4_manifest, "_checkout_snapshot", side_effect=snapshot + ), mock.patch.object( + gemma4_manifest, + "_derive_owned_paths", + return_value=[self.logical_path], + ): + return gemma4_manifest.create_source_manifest( + self.fbsource_root, + self.oss_root, + ) + + def test_source_manifest_producer_round_trip(self) -> None: + manifest = self._create_source_manifest() + gemma4_manifest.validate_source_manifest(manifest) + self.assertEqual([entry["path"] for entry in manifest["files"]], [self.logical_path]) + + def test_owned_union_uses_the_reviewed_plain_summaries(self) -> None: + summaries = gemma4_manifest._GEMMA_PRODUCTION_DIFF_SUMMARIES + self.assertEqual( + summaries, + ( + "[ExecuTorch][WebGPU] Add shared model runtime prerequisites", + "[ExecuTorch][Vulkan] Support scoped Gemma symbolic partitioning", + "[ExecuTorch][WebGPU] Add Gemma 4 plain runtime and guarded routes", + "[ExecuTorch][WebGPU] Add Gemma 4 plain export and artifact contract", + "[ExecuTorch][WebGPU] Add plain Gemma 4 source-closure tests", + ), + ) + expected_reverse = tuple(reversed(summaries)) + + def source_control(argv: list[str], _label: str) -> str: + if "log" in argv: + revision = argv[argv.index("-r") + 1] + offset = 0 if revision == "." else int(revision.removeprefix(".~")) + return f"{offset + 1:040x}\n{expected_reverse[offset]}\n" + node = argv[argv.index("--change") + 1] + offset = int(node, 16) - 1 + path = f"runtime/plain_owned_{offset}.cpp" + return f"xplat/executorch/{path}\nfbcode/executorch/{path}\n" + + with mock.patch.object( + gemma4_manifest, "_run_source_control", side_effect=source_control + ): + paths = gemma4_manifest._derive_owned_paths( + self.fbsource_root, summaries + ) + self.assertEqual( + paths, + [f"runtime/plain_owned_{index}.cpp" for index in range(5)], + ) + + def test_create_source_manifest_cli_round_trip(self) -> None: + output = self.root / "source.json" + manifest = _test_source_manifest() + with mock.patch.object( + gemma4_manifest, "create_source_manifest", return_value=manifest + ): + self.assertEqual( + gemma4_manifest.main( + [ + "create-source-manifest", + "--fbsource-root", + str(self.fbsource_root), + "--oss-root", + str(self.oss_root), + "--output", + str(output), + ] + ), + 0, + ) + document = json.loads(output.read_text(encoding="utf-8")) + gemma4_manifest.validate_source_manifest(document) + + def test_create_wgsl_manifest_cli_round_trip(self) -> None: + output = self.root / "wgsl.json" + manifest = _test_wgsl_manifest() + with mock.patch.object( + gemma4_manifest, "create_wgsl_manifest", return_value=manifest + ): + self.assertEqual( + gemma4_manifest.main( + [ + "create-wgsl-manifest", + "--backend-root", + str(self.fbsource_root), + "--output", + str(output), + ] + ), + 0, + ) + document = json.loads(output.read_text(encoding="utf-8")) + gemma4_manifest.validate_wgsl_manifest(document) + + def test_create_source_receipt_cli_round_trip(self) -> None: + output = self.root / "receipt.json" + receipt = _sealed_source_receipt() + with mock.patch.object( + gemma4_manifest, "create_source_closure_receipt", return_value=receipt + ): + self.assertEqual( + gemma4_manifest.main( + [ + "create-source-receipt", + "--fbsource-root", + str(self.fbsource_root), + "--oss-root", + str(self.oss_root), + "--backend-root", + str(self.fbsource_root), + "--output", + str(output), + ] + ), + 0, + ) + self.assertEqual(json.loads(output.read_text(encoding="utf-8")), receipt) + + def test_wgsl_producer_rejects_stale_generated_output(self) -> None: + backend_root = self.root / "fbsource/xplat/executorch/backends/webgpu" + shader = backend_root / "runtime/ops/add/binary_add.wgsl" + header = backend_root / "runtime/ops/add/binary_add_wgsl.h" + registry = backend_root / "runtime/WebGPUShaderRegistry.cpp" + generator_path = backend_root / "scripts/gen_wgsl_headers.py" + for path, contents in ( + (shader, b"shader"), + (header, b"stale"), + (registry, b"registry"), + (generator_path, b"generator"), + ): + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(contents) + + class Generator: + def discover(self) -> list[Path]: + return [shader] + + def collect_outputs(self) -> tuple[dict[Path, bytes], list[Path]]: + return {header: b"fresh", registry: b"registry"}, [] + + def registry_path(self) -> Path: + return registry + + with mock.patch.object( + gemma4_manifest, + "_checkout_snapshot", + return_value={"clean": True, "head": "1" * 40}, + ), mock.patch.object( + gemma4_manifest, "_load_wgsl_generator", return_value=Generator() + ): + with self.assertRaisesRegex(ValueError, "generated output is stale"): + gemma4_manifest.create_wgsl_manifest(backend_root)