diff --git a/backends/webgpu/CMakeLists.txt b/backends/webgpu/CMakeLists.txt index 919a849a7a6..c02c1032a13 100644 --- a/backends/webgpu/CMakeLists.txt +++ b/backends/webgpu/CMakeLists.txt @@ -327,6 +327,12 @@ if(EXECUTORCH_BUILD_WEBGPU_TEST) target_link_libraries(webgpu_dynamic_shape_test PRIVATE GTest::gtest) add_webgpu_native_test(webgpu_index_test test/native/test_index.cpp) target_link_libraries(webgpu_index_test PRIVATE GTest::gtest) + add_webgpu_native_test(webgpu_topk_test test/native/test_topk.cpp) + target_link_libraries(webgpu_topk_test PRIVATE GTest::gtest) + add_webgpu_native_test(webgpu_scatter_test test/native/test_scatter.cpp) + target_link_libraries(webgpu_scatter_test PRIVATE GTest::gtest) + add_webgpu_native_test(webgpu_q4gsw_m3_test test/native/test_q4gsw_m3.cpp) + target_link_libraries(webgpu_q4gsw_m3_test PRIVATE GTest::gtest) # Device-free fold unit test (gtest_main provides main; no device needed). add_webgpu_native_test( diff --git a/backends/webgpu/scripts/test_webgpu_native_ci.sh b/backends/webgpu/scripts/test_webgpu_native_ci.sh index fcf81066cdc..15910f30928 100644 --- a/backends/webgpu/scripts/test_webgpu_native_ci.sh +++ b/backends/webgpu/scripts/test_webgpu_native_ci.sh @@ -35,12 +35,21 @@ fi cd "${EXECUTORCH_ROOT}" +bash "${SCRIPT_DIR}/test_gemma4_wasm_factory_contract.sh" --validate-names +buck2 test fbcode//executorch/backends/webgpu/test:test_wgsl_codegen + # ── Exports for the model-driven executables ───────────────────────────────── if ! "${PYTHON_EXECUTABLE}" -c "import executorch" 2>/dev/null; then echo "ERROR: executorch wheel unavailable; required fixture exports cannot run" >&2 exit 1 fi +# ── Source contracts: no Buck target, they read across packages ────────────── +$PYTHON_EXECUTABLE -m unittest \ + executorch.backends.webgpu.test.test_native_ci_contract +$PYTHON_EXECUTABLE -m unittest \ + executorch.examples.models.gemma4.tests.test_oss_source_closure + require_file() { if [[ ! -f "$1" ]]; then echo "ERROR: required WebGPU fixture missing: $1" >&2 @@ -48,6 +57,17 @@ require_file() { fi } +recreate_exact_directory() { + local target="$1" + local expected="$2" + if [[ "${target}" != "${expected}" ]]; then + echo "ERROR: refusing to recreate unexpected directory: ${target}" >&2 + return 1 + fi + rm -rf -- "${target}" + mkdir -p -- "${target}" +} + run_with_required_device() { local output if ! output="$("$@" 2>&1)"; then @@ -84,7 +104,7 @@ run_required_gtests() { echo "ERROR: required WebGPU run did not pass exactly three tests" >&2 return 1 fi - if grep -Eq '^\\[ SKIPPED \\]' <<<"${output}"; then + if grep -Eq '^\[ SKIPPED \]' <<<"${output}"; then echo "ERROR: required WebGPU run skipped a test" >&2 return 1 fi @@ -93,6 +113,9 @@ run_required_gtests() { DISPATCH_ORDER_DIR="/tmp/dispatch_order" UPDATE_CACHE_DIR="/tmp/update_cache" INDEX_DIR="/tmp/index" +TOPK_DIR="/tmp/topk" +TOPK_AUTHORITY="/tmp/topk_eager_authority.json" +SCATTER_DIR="/tmp/scatter" DYNAMIC_SHAPE_DIR="/tmp/dynamic_shape" ROPE_HF_DIR="/tmp/webgpu_rope_hf" SYMINT_BLOB="/tmp/sdpa_dyn_small.pte" @@ -116,6 +139,8 @@ PREPACK2_GOLDEN="/tmp/webgpu_prepack_two_const_golden.bin" PREPACK_TIED_MODEL="/tmp/webgpu_prepack_tied_const.pte" PREPACK_TIED_GOLDEN="/tmp/webgpu_prepack_tied_const_golden.bin" +recreate_exact_directory "${UPDATE_CACHE_DIR}" "/tmp/update_cache" + $PYTHON_EXECUTABLE -c " from executorch.backends.webgpu.test.ops.test_quantized_linear import export_all_quantized_linear_models, export_output_suppression_models export_all_quantized_linear_models('/tmp') @@ -153,6 +178,8 @@ export_dispatch_order_cases('${DISPATCH_ORDER_DIR}') $PYTHON_EXECUTABLE -c " from executorch.backends.webgpu.test.ops.test_update_cache import ( + export_dynamic_update_cache, + export_intermediate_dynamic_update_cache, export_update_cache_cases, export_update_cache_replay, export_update_cache_negative, @@ -160,18 +187,38 @@ from executorch.backends.webgpu.test.ops.test_update_cache import ( export_update_cache_cases('${UPDATE_CACHE_DIR}') export_update_cache_replay('${UPDATE_CACHE_DIR}') export_update_cache_negative('${UPDATE_CACHE_DIR}') +export_dynamic_update_cache('${UPDATE_CACHE_DIR}/dynamic.pte') +export_intermediate_dynamic_update_cache('${UPDATE_CACHE_DIR}/dynamic_intermediate.pte') " +require_file "${UPDATE_CACHE_DIR}/dynamic.pte" +require_file "${UPDATE_CACHE_DIR}/dynamic_intermediate.pte" $PYTHON_EXECUTABLE -c " from executorch.backends.webgpu.test.ops.index.test_index import export_all_index_models export_all_index_models('${INDEX_DIR}') " +# The exporter validates this receipt against the committed authority digest. +EAGLE_TOPK_EAGER_RECEIPT="${TOPK_AUTHORITY}" $PYTHON_EXECUTABLE -m unittest \ + executorch.backends.webgpu.test.ops.topk.test_topk.TestEagleTopKCpu.test_eager_reference_is_repeatable + +$PYTHON_EXECUTABLE -m executorch.backends.webgpu.test.ops.topk.export_topk_artifacts \ + "${TOPK_DIR}" "${TOPK_AUTHORITY}" + +$PYTHON_EXECUTABLE -m executorch.backends.webgpu.test.ops.scatter.export_scatter_artifacts \ + "${SCATTER_DIR}" + +recreate_exact_directory "${DYNAMIC_SHAPE_DIR}" "/tmp/dynamic_shape" 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" +require_file "${DYNAMIC_SHAPE_DIR}/dyn_slice_2d.pte" +require_file "${DYNAMIC_SHAPE_DIR}/slice_dual_store.pte" +require_file "${DYNAMIC_SHAPE_DIR}/slice_dual_store.input.bin" +require_file "${DYNAMIC_SHAPE_DIR}/slice_dual_store.out0.golden.bin" +require_file "${DYNAMIC_SHAPE_DIR}/slice_dual_store.out1.golden.bin" $PYTHON_EXECUTABLE -c " from executorch.backends.webgpu.test.ops.test_sdpa import ( @@ -189,6 +236,10 @@ export_incache_decode('/tmp') require_file "${ROPE_HF_DIR}/rope_hf_dynamic.pte" require_file "${SYMINT_BLOB}" require_file "${OUTPUT_SUPPRESSION_DIR}/input.bin" +require_file "${TOPK_AUTHORITY}" +require_file "${TOPK_DIR}/cases.txt" +require_file "${SCATTER_DIR}/cases.txt" +require_file "${SCATTER_DIR}/base.bin" # ── Configure (Dawn-only: no -DWEBGPU_IMPL; Dawn is the sole backend) ───────── echo "=== Configure WebGPU native tests on Dawn ===" @@ -209,7 +260,7 @@ cmake \ "${EXECUTORCH_ROOT}" # ── Build + run every fixed native test target in this tree ────────────────── -REQUIRED_TARGETS=(webgpu_native_test webgpu_dispatch_order_test webgpu_scratch_buffer_test webgpu_update_cache_test webgpu_index_test webgpu_dynamic_shape_test webgpu_dispatch_2d_test webgpu_compute_dispatch_test webgpu_execution_options_test webgpu_output_suppression_test webgpu_op_test_util_test) +REQUIRED_TARGETS=(webgpu_native_test webgpu_dispatch_order_test webgpu_scratch_buffer_test webgpu_update_cache_test webgpu_update_cache_state_test webgpu_index_test webgpu_dynamic_shape_test webgpu_dispatch_2d_test webgpu_compute_dispatch_test webgpu_execution_options_test webgpu_output_suppression_test webgpu_op_test_util_test webgpu_topk_test webgpu_scatter_test webgpu_q4gsw_m3_test) BIN_DIR="${BUILD_DIR}/backends/webgpu" DEFINED_TARGETS="$(cmake --build "${BUILD_DIR}" --target help 2>/dev/null || true)" @@ -224,6 +275,7 @@ for t in "${REQUIRED_TARGETS[@]}"; do done echo "=== Run native tests on Dawn + SwiftShader ===" +"${BIN_DIR}/webgpu_update_cache_state_test" run_with_required_device env WEBGPU_TEST_SDPA_DIR=/tmp/ \ WEBGPU_TEST_QUANTIZED_LINEAR_DIR=/tmp/ \ WEBGPU_TEST_EMBEDDING_Q4GSW_MODEL="${EMBEDDING_MODEL}" \ @@ -247,9 +299,14 @@ run_with_required_device env WEBGPU_TEST_SDPA_DIR=/tmp/ \ WEBGPU_TEST_PREPACK_TIED_MODEL="${PREPACK_TIED_MODEL}" \ WEBGPU_TEST_PREPACK_TIED_GOLDEN="${PREPACK_TIED_GOLDEN}" \ "${BIN_DIR}/webgpu_native_test" -"${BIN_DIR}/webgpu_update_cache_test" "${UPDATE_CACHE_DIR}" +run_with_required_device env WEBGPU_REQUIRE_DEVICE=1 \ + WEBGPU_UPDATE_CACHE_DIR="${UPDATE_CACHE_DIR}" \ + "${BIN_DIR}/webgpu_update_cache_test" "${UPDATE_CACHE_DIR}" "${BIN_DIR}/webgpu_dispatch_order_test" "${DISPATCH_ORDER_DIR}" "${BIN_DIR}/webgpu_index_test" "${INDEX_DIR}" +"${BIN_DIR}/webgpu_topk_test" "${TOPK_DIR}" +"${BIN_DIR}/webgpu_scatter_test" "${SCATTER_DIR}" +"${BIN_DIR}/webgpu_q4gsw_m3_test" "${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}" \ diff --git a/backends/webgpu/test/BUCK b/backends/webgpu/test/BUCK index 079fb77cb81..ef51fd7477f 100644 --- a/backends/webgpu/test/BUCK +++ b/backends/webgpu/test/BUCK @@ -57,6 +57,90 @@ fbcode_target( ], ) +fbcode_target( + _kind = python_unittest, + name = "test_to_copy", + srcs = [ + "ops/test_to_copy.py", + ], + deps = [ + "//caffe2:torch", + "//executorch/backends/vulkan/partitioner:vulkan_partitioner", + "//executorch/backends/vulkan:vulkan_preprocess", + "//executorch/exir:lib", + ], +) + +fbcode_target( + _kind = python_unittest, + name = "test_gather", + srcs = [ + "ops/test_gather.py", + ], + deps = [ + "//caffe2:torch", + "//executorch/backends/vulkan/partitioner:vulkan_partitioner", + "//executorch/backends/vulkan:vulkan_preprocess", + "//executorch/exir:lib", + ], +) + +fbcode_target( + _kind = python_unittest, + name = "test_where", + srcs = [ + "ops/test_where.py", + ], + deps = [ + "//caffe2:torch", + "//executorch/backends/vulkan/partitioner:vulkan_partitioner", + "//executorch/backends/vulkan:vulkan_preprocess", + "//executorch/exir:lib", + ], +) + +# Pure-stdlib CPU authority for the top-k route; no torch, no GPU. +fbcode_target( + _kind = python_unittest, + name = "test_topk_cpu", + srcs = [ + "ops/topk/test_topk.py", + ], +) + +# Pure-stdlib CPU authority for the scatter route; no torch, no GPU. +fbcode_target( + _kind = python_unittest, + name = "test_scatter_cpu", + srcs = [ + "ops/scatter/test_scatter.py", + ], +) + +# Source contract over the op-test driver; reads its inputs, so they ship as resources. +fbcode_target( + _kind = python_unittest, + name = "test_typed_input_contract", + srcs = [ + "op_tests/test_typed_input_contract.py", + ], + resources = [ + "op_tests/cases.py", + "op_tests/generate_op_tests.py", + "op_tests/op_test_driver.cpp", + "ops/test_to_copy.py", + "ops/test_where.py", + ], +) + +# The native CI contract reads four Buck packages; test_webgpu_native_ci.sh runs it. +fbcode_target( + _kind = runtime.python_library, + name = "test_native_ci_contract", + srcs = ["test_native_ci_contract.py"], + typing = True, +) + fbcode_target( _kind = python_unittest, name = "test_rope_hf_single", diff --git a/backends/webgpu/test/native/test_q4gsw_m3.cpp b/backends/webgpu/test/native/test_q4gsw_m3.cpp new file mode 100644 index 00000000000..678992cd572 --- /dev/null +++ b/backends/webgpu/test/native/test_q4gsw_m3.cpp @@ -0,0 +1,905 @@ +/* + * 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. + */ + +// Device-graph tests for the linear_q4gsw M==3 shared-bicol route: structural +// route selection, the generic fallback, raw-fp32 scale fidelity (no bf16 +// rounding), resize re-entry across the M==3 boundary, and fail-closed +// validation. Inputs are generated in-process, so no fixture directory. + +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace executorch::backends::webgpu { +namespace { + +// QuantizedLinear.cpp:60-62 -- the M=3 route's workgroup size and the shared +// storage it needs (6 f32 partials per lane). +constexpr uint32_t kM3Invocations = 64u; +constexpr uint32_t kM3StorageBytes = 6u * kM3Invocations * sizeof(float); +constexpr uint32_t kM3PartialArrays = 6u; + +// WebGPU spec minimums (https://www.w3.org/TR/webgpu/#limits); every +// conformant device clears the M=3 gate, so a failure here is a real defect. +constexpr uint32_t kSpecMinInvocations = 256u; +constexpr uint32_t kSpecMinWorkgroupSizeX = 256u; +constexpr uint32_t kSpecMinWorkgroupStorage = 16384u; + +constexpr const char* kQ4gswOp = "et_vk.linear_q4gsw.default"; +constexpr const char* kM3Shader = "q4gsw_linear_m3_shared_bicol"; +constexpr const char* kM3Kernel = "linear_q4gsw_m3_shared_bicol"; +constexpr const char* kBicolKernel = "linear_q4gsw_coop4_bicol"; +constexpr const char* kTiledKernel = "linear_q4gsw_tiled"; + +// K % 8 == 0 and group_size % 8 == 0 make the shape bicol/M3 eligible +// (QuantizedLinear.cpp:492-495); K % 16 != 0 keeps every non-M3 case on the +// fp32 tiled kernel (steel_workgroup_count, QuantizedLinear.cpp:133-143) and +// K/N stay under the shmem thresholds (QuantizedLinear.cpp:52-53). +constexpr int64_t kK = 72; +constexpr int64_t kGroupSize = 8; +constexpr int64_t kGroups = kK / kGroupSize; +constexpr int64_t kKPacked = kK / 2; +constexpr int64_t kN = 6; +constexpr int64_t kMaxM = 4; + +// Project tolerance for cross-kernel fp32 comparisons; the routes reduce K in +// different orders, so bit-exactness is not a contract between them. +constexpr float kAtol = 1e-3f; +constexpr float kRtol = 1e-3f; +// Scale-fidelity gate: >=19x tighter than the >=1.95e-3 relative shift any +// bf16 rounding of the fixture's scales would cause. +constexpr double kScaleGate = 1e-4; +constexpr double kBf16Separation = 1e-3; + +// NOLINTBEGIN(cppcoreguidelines-avoid-non-const-global-variables) +WGPUDevice g_device = nullptr; +// NOLINTEND(cppcoreguidelines-avoid-non-const-global-variables) + +// Mirrors m3_shared_supported (QuantizedLinear.cpp:103-109). +bool m3_shared_supported_here(WGPUDevice device) { + WGPULimits limits = {}; + return wgpuDeviceGetLimits(device, &limits) == WGPUStatus_Success && + limits.maxComputeInvocationsPerWorkgroup >= kM3Invocations && + limits.maxComputeWorkgroupSizeX >= kM3Invocations && + limits.maxComputeWorkgroupStorageSize >= kM3StorageBytes; +} + +struct Q4gswSpec { + std::vector input_dims = {3u, static_cast(kK)}; + uint32_t n = static_cast(kN); + uint32_t k_packed = static_cast(kKPacked); + uint32_t groups = static_cast(kGroups); + uint32_t padded_n = static_cast(kN); + int64_t group_size = kGroupSize; + uint32_t bias_elems = 0u; // 0 -> Null bias arg + bool dynamic = false; // emit sym_size.int so the input carries dynamic dims + bool fp16_scales = false; + bool rank1_scales = false; + bool null_group_size = false; + int id_shift = 0; // unused leading values shift every value id + int mem_obj_base = 0; +}; + +struct Q4gswIds { + int input = 0; + int weight = 0; + int scales = 0; + int bias = 0; + int output = 0; +}; + +Q4gswIds build_q4gsw_graph(WebGPUGraph& graph, const Q4gswSpec& spec) { + namespace vk = vkgraph; + ::flatbuffers::FlatBufferBuilder fbb; + std::vector<::flatbuffers::Offset> values; + auto add_tensor = [&](vk::VkDataType dtype, + const std::vector& dims, + int mem_obj_id) { + const int id = static_cast(values.size()); + values.push_back(vk::CreateVkValue( + fbb, + vk::GraphTypes::VkTensor, + vk::CreateVkTensorDirect( + fbb, dtype, &dims, /*constant_id=*/-1, mem_obj_id) + .Union())); + return id; + }; + auto add_int = [&](int64_t value) { + const int id = static_cast(values.size()); + values.push_back(vk::CreateVkValue( + fbb, vk::GraphTypes::Int, vk::CreateInt(fbb, value).Union())); + return id; + }; + auto add_null = [&]() { + const int id = static_cast(values.size()); + values.push_back(vk::CreateVkValue(fbb)); + return id; + }; + + for (int i = 0; i < spec.id_shift; i++) { + add_int(i); + } + + std::vector output_dims = spec.input_dims; + output_dims.back() = spec.n; + + Q4gswIds ids; + ids.input = + add_tensor(vk::VkDataType::FLOAT32, spec.input_dims, spec.mem_obj_base); + ids.weight = add_tensor( + vk::VkDataType::UINT8, {spec.n, spec.k_packed}, spec.mem_obj_base + 1); + const vk::VkDataType scales_dtype = + spec.fp16_scales ? vk::VkDataType::FLOAT16 : vk::VkDataType::FLOAT32; + ids.scales = spec.rank1_scales + ? add_tensor(scales_dtype, {spec.groups}, spec.mem_obj_base + 2) + : add_tensor( + scales_dtype, {spec.groups, spec.padded_n}, spec.mem_obj_base + 2); + const int group_size = + spec.null_group_size ? add_null() : add_int(spec.group_size); + ids.bias = spec.bias_elems == 0u + ? add_null() + : add_tensor( + vk::VkDataType::FLOAT32, {spec.bias_elems}, spec.mem_obj_base + 4); + ids.output = + add_tensor(vk::VkDataType::FLOAT32, output_dims, spec.mem_obj_base + 3); + + std::vector<::flatbuffers::Offset> chain; + if (spec.dynamic) { + const int dim = add_int(0); + const int symint = static_cast(values.size()); + values.push_back(vk::CreateVkValue( + fbb, vk::GraphTypes::SymInt, vk::CreateSymInt(fbb, 0).Union())); + const std::vector sym_args = {ids.input, dim, symint}; + chain.push_back( + vk::CreateOperatorCallDirect(fbb, 0, "sym_size.int", &sym_args)); + } + const std::vector q4_args = { + ids.input, ids.weight, ids.scales, group_size, ids.bias, ids.output}; + chain.push_back(vk::CreateOperatorCallDirect( + fbb, static_cast(chain.size()), kQ4gswOp, &q4_args)); + + std::vector input_ids = { + static_cast(ids.input), + static_cast(ids.weight), + static_cast(ids.scales)}; + if (spec.bias_elems != 0u) { + input_ids.push_back(static_cast(ids.bias)); + } + const std::vector output_ids = {static_cast(ids.output)}; + const auto root = vk::CreateVkGraphDirect( + fbb, "0", &chain, &values, &input_ids, &output_ids); + vk::FinishVkGraphBuffer(fbb, root); + + graph.set_device(g_device); + graph.build(fbb.GetBufferPointer(), nullptr, 0, nullptr); + return ids; +} + +struct HostFixture { + std::vector input; // rows * K, row-major + std::vector weight; // N * K_packed + std::vector scales; // groups * padded_N + std::vector bias; // N, empty when the graph has no bias +}; + +uint32_t next_u32(uint32_t& state) { + state = state * 1664525u + 1013904223u; + return state; +} + +float unit_float(uint32_t& state) { + return static_cast(next_u32(state) >> 8u) / 8388608.0f - 1.0f; +} + +// A value exactly halfway between two bf16 neighbours (mantissa low bits +// 0x8000), so ANY bf16 rounding of it moves the value >= 2^-9 relative. +float bf16_midpoint(uint32_t index, int exponent) { + const uint32_t bits = (static_cast(127 + exponent) << 23u) | + ((index & 0x7Fu) << 16u) | 0x8000u; + float value = 0.0f; + std::memcpy(&value, &bits, sizeof(value)); + return value; +} + +float truncate_to_bf16(float value) { + uint32_t bits = 0u; + std::memcpy(&bits, &value, sizeof(bits)); + bits &= 0xFFFF0000u; + float truncated = 0.0f; + std::memcpy(&truncated, &bits, sizeof(truncated)); + return truncated; +} + +HostFixture make_random_fixture(int64_t rows, bool with_bias) { + HostFixture fixture; + uint32_t state = 0x5eed1234u; + fixture.input.resize(static_cast(rows * kK)); + for (float& value : fixture.input) { + value = unit_float(state); + } + fixture.weight.resize(static_cast(kN * kKPacked)); + for (uint8_t& byte : fixture.weight) { + byte = static_cast(next_u32(state) >> 24u); + } + fixture.scales.resize(static_cast(kGroups * kN)); + for (size_t i = 0; i < fixture.scales.size(); i++) { + fixture.scales[i] = + bf16_midpoint(static_cast(11u * i + 3u), /*exponent=*/-6); + } + if (with_bias) { + fixture.bias.resize(static_cast(kN)); + for (float& value : fixture.bias) { + value = unit_float(state); + } + } + return fixture; +} + +// All nibbles are 0xF (dequant +7, q4gsw_linear.wgsl:69) and row r is the +// constant r+1, so out[r][c] == 7 * (r+1) * group_size * sum_g scales[g][c]. +HostFixture make_analytic_fixture(int64_t rows) { + HostFixture fixture; + fixture.input.resize(static_cast(rows * kK)); + for (int64_t r = 0; r < rows; r++) { + for (int64_t k = 0; k < kK; k++) { + fixture.input[static_cast(r * kK + k)] = + static_cast(r + 1); + } + } + fixture.weight.assign(static_cast(kN * kKPacked), 0xFFu); + fixture.scales.resize(static_cast(kGroups * kN)); + for (size_t i = 0; i < fixture.scales.size(); i++) { + fixture.scales[i] = + bf16_midpoint(static_cast(17u * i + 5u), /*exponent=*/0); + } + return fixture; +} + +double +analytic_expected(const std::vector& scales, int64_t row, int64_t col) { + double scale_sum = 0.0; + for (int64_t g = 0; g < kGroups; g++) { + scale_sum += scales[static_cast(g * kN + col)]; + } + return 7.0 * static_cast(row + 1) * static_cast(kGroupSize) * + scale_sum; +} + +std::vector +run_graph(WebGPUGraph& graph, const HostFixture& fixture, int64_t m) { + std::vector inputs; + inputs.push_back( + {fixture.input.data(), + static_cast(m * kK) * sizeof(float), + false, + true}); + inputs.push_back( + {fixture.weight.data(), fixture.weight.size(), false, false}); + inputs.push_back( + {fixture.scales.data(), + fixture.scales.size() * sizeof(float), + false, + true}); + if (!fixture.bias.empty()) { + inputs.push_back( + {fixture.bias.data(), + fixture.bias.size() * sizeof(float), + false, + true}); + } + std::vector out(static_cast(m * kN), 0.0f); + std::vector outputs(1); + outputs[0] = {out.data(), out.size() * sizeof(float), /*host_is_fp32=*/true}; + + graph.copy_inputs(inputs); + const WebGPUExecutionPlan plan = graph.make_execution_plan({}); + graph.execute(plan); + graph.copy_outputs(outputs, plan); + return out; +} + +struct ActiveDispatch { + std::string kernel_name; + uint32_t workgroup_count_x = 0u; + size_t count = 0u; +}; + +// The dispatches a route group left runnable; a selected route has exactly one. +ActiveDispatch active_dispatch(WebGPUGraph& graph) { + ActiveDispatch active; + for (size_t i = 0; i < graph.num_dispatches(); i++) { + const WebGPUDispatch& dispatch = graph.dispatch_at(i); + if (dispatch.kind != WebGPUDispatch::Kind::Compute || + dispatch.workgroup_count_x == 0u) { + continue; + } + active.kernel_name = dispatch.kernel_name; + active.workgroup_count_x = dispatch.workgroup_count_x; + active.count++; + } + return active; +} + +size_t count_kernel(WebGPUGraph& graph, const char* kernel_name) { + size_t count = 0; + for (size_t i = 0; i < graph.num_dispatches(); i++) { + if (graph.dispatch_at(i).kernel_name == kernel_name) { + count++; + } + } + return count; +} + +void expect_close( + const std::vector& got, + const std::vector& want, + size_t count, + const std::string& label) { + ASSERT_GE(got.size(), count) << label; + ASSERT_GE(want.size(), count) << label; + for (size_t i = 0; i < count; i++) { + ASSERT_TRUE(std::isfinite(got[i])) << label << " i=" << i; + const float abs_err = std::fabs(got[i] - want[i]); + const float rel_err = abs_err / std::fmax(std::fabs(want[i]), 1e-6f); + EXPECT_TRUE(abs_err <= kAtol || rel_err <= kRtol) + << label << " i=" << i << " got=" << got[i] << " want=" << want[i]; + } +} + +void expect_build_error(const Q4gswSpec& spec, const char* expected) { + WebGPUGraph graph; + std::string error; + try { + build_q4gsw_graph(graph, spec); + } catch (const std::exception& exception) { + error = exception.what(); + } + EXPECT_EQ(error, expected); + EXPECT_EQ(graph.memory_stats().num_dispatches, 0); +} + +// Strips `//` and `/* */` so a shader's prose cannot satisfy or trip a check. +std::string strip_wgsl_comments(const std::string& src) { + std::string out; + out.reserve(src.size()); + for (size_t i = 0; i < src.size();) { + if (src.compare(i, 2, "//") == 0) { + while (i < src.size() && src[i] != '\n') { + i++; + } + } else if (src.compare(i, 2, "/*") == 0) { + i += 2; + while (i + 1 < src.size() && src.compare(i, 2, "*/") != 0) { + i++; + } + i = i + 1 < src.size() ? i + 2 : src.size(); + } else { + out.push_back(src[i++]); + } + } + return out; +} + +// The declared x dim, resolving a `const : u32 = u;` indirection. +uint32_t declared_workgroup_size_x(const std::string& src) { + const std::string code = strip_wgsl_comments(src); + const size_t at = code.find("@workgroup_size"); + if (at == std::string::npos) { + return 0; + } + const size_t open = code.find('(', at); + const size_t close = code.find(')', open); + if (open == std::string::npos || close == std::string::npos) { + return 0; + } + std::string token = code.substr(open + 1, close - open - 1); + token = token.substr(0, token.find(',')); + const size_t begin = token.find_first_not_of(" \t\n\r"); + if (begin == std::string::npos) { + return 0; + } + token = token.substr(begin, token.find_last_not_of(" \t\n\r") + 1 - begin); + if (!token.empty() && std::isdigit(static_cast(token[0]))) { + return static_cast(std::strtoul(token.c_str(), nullptr, 10)); + } + const size_t decl = code.find("const " + token); + const size_t eq = + decl == std::string::npos ? std::string::npos : code.find('=', decl); + if (eq == std::string::npos) { + return 0; + } + return static_cast( + std::strtoul(code.c_str() + eq + 1, nullptr, 10)); +} + +std::string to_lower(const std::string& text) { + std::string out = text; + for (char& c : out) { + c = static_cast(std::tolower(static_cast(c))); + } + return out; +} + +size_t count_occurrences(const std::string& haystack, const char* needle) { + size_t count = 0; + const std::string pattern(needle); + for (size_t at = haystack.find(pattern); at != std::string::npos; + at = haystack.find(pattern, at + pattern.size())) { + count++; + } + return count; +} + +} // namespace + +// The M=3 route is gated on device limits, not skipped: every WebGPU-conformant +// device clears them, so an unmet gate is a defect, not an unsupported setup. +TEST(Q4gswM3Device, ClearsTheSharedBicolLimitGate) { + WGPULimits limits = {}; + ASSERT_EQ(wgpuDeviceGetLimits(g_device, &limits), WGPUStatus_Success); + EXPECT_GE(limits.maxComputeInvocationsPerWorkgroup, kSpecMinInvocations); + EXPECT_GE(limits.maxComputeWorkgroupSizeX, kSpecMinWorkgroupSizeX); + EXPECT_GE(limits.maxComputeWorkgroupStorageSize, kSpecMinWorkgroupStorage); + ASSERT_TRUE(m3_shared_supported_here(g_device)) + << "device is below the M=3 gate (64 invocations, 64 x, " + << kM3StorageBytes << " shared bytes)"; +} + +// The op compiles the M=3 pipeline from kQ4gswLinearM3SharedBicolWGSL directly +// (QuantizedLinear.cpp:635, :694), so pin the registry to that same source and +// assert the source itself carries no scale truncation. +TEST(Q4gswM3Shader, ServesRawFp32ScalesWithNoBf16Rounding) { + const WebGPUShaderInfo& info = get_webgpu_shader_info(kM3Shader); + ASSERT_NE(info.source, nullptr); + EXPECT_STREQ(info.source, kQ4gswLinearM3SharedBicolWGSL); + EXPECT_EQ(declared_workgroup_size_x(info.source), kM3Invocations); + EXPECT_EQ(info.workgroup_size_x, kM3Invocations) + << "registry workgroup size (kQ4gswLinearM3SharedBicolWorkgroupSizeX) " + "disagrees with @workgroup_size in the shader source"; + + const std::string src = strip_wgsl_comments(info.source); + const std::string lowered = to_lower(src); + EXPECT_NE( + src.find("var t_scales: array"), std::string::npos) + << "t_scales must stay a raw fp32 storage array"; + for (const char* needle : {"f16", "bfloat", "0xffff0000", "bitcast"}) { + EXPECT_EQ(lowered.find(needle), std::string::npos) + << "scale-truncation pattern '" << needle << "' reappeared"; + } + // A 7th partial array would need more shared memory than the CPU gate asks + // the device for (kQ4gswM3StorageBytes == 6 * 64 * 4). + EXPECT_EQ(count_occurrences(src, "var"), kM3PartialArrays); + EXPECT_EQ(count_occurrences(src, "array"), kM3PartialArrays); +} + +TEST(Q4gswM3Route, SelectsSharedBicolOnlyAtMEqualsThree) { + const struct { + int64_t m; + const char* kernel; + } cases[] = { + {1, kBicolKernel}, + {2, kTiledKernel}, + {3, kM3Kernel}, + {4, kTiledKernel}, + }; + for (const auto& test_case : cases) { + SCOPED_TRACE(std::string("M=") + std::to_string(test_case.m)); + Q4gswSpec spec; + spec.input_dims = { + static_cast(test_case.m), static_cast(kK)}; + WebGPUGraph graph; + ASSERT_NO_THROW(build_q4gsw_graph(graph, spec)); + ASSERT_EQ(graph.num_dispatches(), 1u); + EXPECT_EQ(graph.dispatch_at(0).kernel_name, test_case.kernel); + EXPECT_EQ(count_kernel(graph, kM3Kernel), test_case.m == 3 ? 1u : 0u); + EXPECT_NE(graph.dispatch_at(0).pipeline, nullptr); + EXPECT_NE(graph.dispatch_at(0).bind_group, nullptr); + EXPECT_GT(graph.dispatch_at(0).workgroup_count_x, 0u); + } +} + +// M=3 dispatches ceil(N/2) column-pair workgroups (QuantizedLinear.cpp:189). +TEST(Q4gswM3Route, DispatchesOneWorkgroupPerColumnPair) { + for (uint32_t n : {6u, 7u}) { + SCOPED_TRACE(std::string("N=") + std::to_string(n)); + Q4gswSpec spec; + spec.n = n; + spec.padded_n = n; + WebGPUGraph graph; + ASSERT_NO_THROW(build_q4gsw_graph(graph, spec)); + ASSERT_EQ(graph.num_dispatches(), 1u); + EXPECT_EQ(graph.dispatch_at(0).kernel_name, kM3Kernel); + EXPECT_EQ(graph.dispatch_at(0).workgroup_count_x, (n + 1u) / 2u); + EXPECT_EQ(graph.dispatch_at(0).workgroup_count_y, 1u); + } +} + +// M comes from numel/K and the device limits alone -- not from a value id, a +// dims pattern, a memory-object id, or any artifact-bound role map. +TEST(Q4gswM3Route, SelectionIsStructuralNotArtifactBound) { + const struct { + const char* name; + Q4gswSpec spec; + } cases[] = { + {"shifted value ids", + [] { + Q4gswSpec s; + s.id_shift = 9; + return s; + }()}, + {"shifted memory-object ids", + [] { + Q4gswSpec s; + s.mem_obj_base = 7; + return s; + }()}, + {"rank-3 input [1,3,K]", + [] { + Q4gswSpec s; + s.input_dims = {1u, 3u, static_cast(kK)}; + return s; + }()}, + {"wider N", + [] { + Q4gswSpec s; + s.n = 10u; + s.padded_n = 10u; + return s; + }()}, + {"padded scales table", + [] { + Q4gswSpec s; + s.padded_n = 16u; + return s; + }()}, + {"biased", + [] { + Q4gswSpec s; + s.bias_elems = static_cast(kN); + return s; + }()}, + }; + for (const auto& test_case : cases) { + SCOPED_TRACE(test_case.name); + WebGPUGraph graph; + ASSERT_NO_THROW(build_q4gsw_graph(graph, test_case.spec)); + ASSERT_EQ(graph.num_dispatches(), 1u); + EXPECT_EQ(graph.dispatch_at(0).kernel_name, kM3Kernel); + EXPECT_EQ(count_kernel(graph, kM3Kernel), 1u); + } +} + +// M3 ineligible (K % 8 or group_size % 8) must still build and dispatch a +// named q4gsw kernel -- never zero dispatches, never a throw. +TEST(Q4gswM3Route, FallsBackToANamedKernelWhenIneligible) { + const struct { + const char* name; + Q4gswSpec spec; + } cases[] = { + {"K % 8 != 0", + [] { + Q4gswSpec s; + s.input_dims = {3u, 20u}; + s.k_packed = 10u; + s.groups = 3u; + return s; + }()}, + {"group_size % 8 != 0", + [] { + Q4gswSpec s; + s.input_dims = {3u, 24u}; + s.k_packed = 12u; + s.groups = 6u; + s.group_size = 4; + return s; + }()}, + }; + for (const auto& test_case : cases) { + SCOPED_TRACE(test_case.name); + WebGPUGraph graph; + ASSERT_NO_THROW(build_q4gsw_graph(graph, test_case.spec)); + ASSERT_EQ(graph.num_dispatches(), 1u); + EXPECT_EQ(graph.dispatch_at(0).kernel_name, kTiledKernel); + EXPECT_EQ(count_kernel(graph, kM3Kernel), 0u); + EXPECT_NE(graph.dispatch_at(0).pipeline, nullptr); + EXPECT_GT(graph.dispatch_at(0).workgroup_count_x, 0u); + } +} + +// The M=4 graph runs the generic tiled kernel over the same weights, scales and +// first three input rows, so its rows 0-2 are a same-math oracle for M=3. +TEST(Q4gswM3Numerics, MatchesTheGenericRouteOnTheSameInputs) { + for (bool with_bias : {false, true}) { + SCOPED_TRACE(with_bias ? "biased" : "unbiased"); + const HostFixture fixture = make_random_fixture(kMaxM, with_bias); + + Q4gswSpec generic_spec; + generic_spec.input_dims = { + static_cast(kMaxM), static_cast(kK)}; + generic_spec.bias_elems = with_bias ? static_cast(kN) : 0u; + WebGPUGraph generic; + ASSERT_NO_THROW(build_q4gsw_graph(generic, generic_spec)); + ASSERT_EQ(generic.num_dispatches(), 1u); + ASSERT_EQ(generic.dispatch_at(0).kernel_name, kTiledKernel); + const std::vector want = run_graph(generic, fixture, kMaxM); + + Q4gswSpec m3_spec; + m3_spec.bias_elems = generic_spec.bias_elems; + WebGPUGraph m3; + ASSERT_NO_THROW(build_q4gsw_graph(m3, m3_spec)); + ASSERT_EQ(m3.num_dispatches(), 1u); + ASSERT_EQ(m3.dispatch_at(0).kernel_name, kM3Kernel); + const std::vector got = run_graph(m3, fixture, 3); + + expect_close(got, want, static_cast(3 * kN), "m3 vs generic"); + } +} + +// Kills a reintroduced bf16 scale rounding numerically: the fixture's scales +// sit at bf16 midpoints, so rounding them (either direction) shifts every +// output by >= 1.95e-3 relative, 19x the gate this asserts. +TEST(Q4gswM3Numerics, AppliesRawFp32ScalesNotBf16) { + const HostFixture fixture = make_analytic_fixture(kMaxM); + for (float scale : fixture.scales) { + const double relative_shift = + std::fabs(scale - truncate_to_bf16(scale)) / scale; + ASSERT_GT(relative_shift, 1.5e-3) + << "fixture scale " << scale << " is not bf16-sensitive"; + } + + Q4gswSpec m3_spec; + WebGPUGraph m3; + ASSERT_NO_THROW(build_q4gsw_graph(m3, m3_spec)); + ASSERT_EQ(m3.num_dispatches(), 1u); + ASSERT_EQ(m3.dispatch_at(0).kernel_name, kM3Kernel); + const std::vector got = run_graph(m3, fixture, 3); + + std::vector bf16_scales = fixture.scales; + for (float& scale : bf16_scales) { + scale = truncate_to_bf16(scale); + } + for (int64_t r = 0; r < 3; r++) { + for (int64_t c = 0; c < kN; c++) { + const double observed = got[static_cast(r * kN + c)]; + const double expected = analytic_expected(fixture.scales, r, c); + const double rounded = analytic_expected(bf16_scales, r, c); + ASSERT_GT(expected, 0.0); + EXPECT_LT(std::fabs(observed - expected) / expected, kScaleGate) + << "raw-fp32 scale mismatch at [" << r << "," << c << "]"; + EXPECT_GT(std::fabs(observed - rounded) / expected, kBf16Separation) + << "output matches the bf16-rounded counterfactual at [" << r << "," + << c << "]"; + } + } +} + +// The generic route is held to the same raw-fp32 scale contract, so a bf16 +// mutant applied to BOTH routes cannot hide behind route parity. +TEST(Q4gswM3Numerics, GenericRouteAlsoAppliesRawFp32Scales) { + const HostFixture fixture = make_analytic_fixture(kMaxM); + Q4gswSpec spec; + spec.input_dims = {static_cast(kMaxM), static_cast(kK)}; + WebGPUGraph graph; + ASSERT_NO_THROW(build_q4gsw_graph(graph, spec)); + ASSERT_EQ(graph.num_dispatches(), 1u); + ASSERT_EQ(graph.dispatch_at(0).kernel_name, kTiledKernel); + const std::vector got = run_graph(graph, fixture, kMaxM); + + for (int64_t r = 0; r < kMaxM; r++) { + for (int64_t c = 0; c < kN; c++) { + const double observed = got[static_cast(r * kN + c)]; + const double expected = analytic_expected(fixture.scales, r, c); + EXPECT_LT(std::fabs(observed - expected) / expected, kScaleGate) + << "generic raw-fp32 scale mismatch at [" << r << "," << c << "]"; + } + } +} + +// One dynamic graph resized across the M==3 boundary in both directions: the +// recorded route must be reselected and stay numerically correct each time. +TEST(Q4gswM3Resize, ReselectsAcrossTheM3Boundary) { + const HostFixture fixture = make_random_fixture(kMaxM, /*with_bias=*/false); + + Q4gswSpec generic_spec; + generic_spec.input_dims = { + static_cast(kMaxM), static_cast(kK)}; + WebGPUGraph generic; + ASSERT_NO_THROW(build_q4gsw_graph(generic, generic_spec)); + ASSERT_EQ(generic.num_dispatches(), 1u); + ASSERT_EQ(generic.dispatch_at(0).kernel_name, kTiledKernel); + const std::vector want = run_graph(generic, fixture, kMaxM); + + Q4gswSpec spec; + spec.input_dims = {static_cast(kMaxM), static_cast(kK)}; + spec.dynamic = true; + WebGPUGraph graph; + Q4gswIds ids; + ASSERT_NO_THROW(ids = build_q4gsw_graph(graph, spec)); + ASSERT_TRUE(graph.has_dynamic_shapes()); + ASSERT_EQ(graph.num_dispatches(), 3u); + EXPECT_EQ(graph.dispatch_at(0).kernel_name, kBicolKernel); + EXPECT_EQ(graph.dispatch_at(1).kernel_name, kM3Kernel); + EXPECT_EQ(graph.dispatch_at(2).kernel_name, kTiledKernel); + EXPECT_EQ(count_kernel(graph, kM3Kernel), 1u); + EXPECT_EQ(active_dispatch(graph).kernel_name, kTiledKernel); + + const struct { + int64_t m; + const char* kernel; + } steps[] = { + {4, kTiledKernel}, + {3, kM3Kernel}, + {1, kBicolKernel}, + {3, kM3Kernel}, + {2, kTiledKernel}, + {4, kTiledKernel}, + }; + for (const auto& step : steps) { + SCOPED_TRACE(std::string("live M=") + std::to_string(step.m)); + graph.resize_input(ids.input, {step.m, kK}); + ASSERT_NO_THROW(graph.propagate_resize()); + const ActiveDispatch active = active_dispatch(graph); + ASSERT_EQ(active.count, 1u); + EXPECT_EQ(active.kernel_name, step.kernel); + if (step.m == 1 || step.m == 3) { + EXPECT_EQ(active.workgroup_count_x, static_cast((kN + 1) / 2)); + } + const std::vector got = run_graph(graph, fixture, step.m); + expect_close( + got, want, static_cast(step.m * kN), "resized vs generic"); + } +} + +TEST(Q4gswM3FailsClosed, RejectsMalformedGraphsBeforeAnyDispatch) { + const struct { + const char* name; + Q4gswSpec spec; + const char* error; + } cases[] = { + {"rank-1 scales", + [] { + Q4gswSpec s; + s.rank1_scales = true; + return s; + }(), + "WebGPU linear_q4gsw: malformed input dims"}, + {"K_packed != ceil(K/2)", + [] { + Q4gswSpec s; + s.k_packed = static_cast(kKPacked) + 1u; + return s; + }(), + "WebGPU linear_q4gsw: K_packed must be ceil(K/2)"}, + {"N*K_packed not u32-aligned", + [] { + Q4gswSpec s; + s.input_dims = {3u, 2u}; + s.n = 1u; + s.k_packed = 1u; + s.groups = 1u; + s.padded_n = 1u; + return s; + }(), + "WebGPU linear_q4gsw: N*K_packed must be a multiple of 4 (u32-packed)"}, + {"fp16 scales", + [] { + Q4gswSpec s; + s.fp16_scales = true; + return s; + }(), + "WebGPU linear_q4gsw: fp32-only (byte-size mismatch)"}, + {"group_size == 0", + [] { + Q4gswSpec s; + s.group_size = 0; + return s; + }(), + "WebGPU linear_q4gsw: group_size <= 0"}, + {"group_size not an Int", + [] { + Q4gswSpec s; + s.null_group_size = true; + return s; + }(), + "WebGPU linear_q4gsw: group_size <= 0"}, + {"too few scale groups", + [] { + Q4gswSpec s; + s.groups = static_cast(kGroups) - 1u; + return s; + }(), + "WebGPU linear_q4gsw: scales dims too small for K/N"}, + {"padded_N < N", + [] { + Q4gswSpec s; + s.padded_n = static_cast(kN) - 1u; + return s; + }(), + "WebGPU linear_q4gsw: scales dims too small for K/N"}, + {"undersized bias", + [] { + Q4gswSpec s; + s.bias_elems = static_cast(kN) - 1u; + return s; + }(), + "WebGPU linear_q4gsw: bias present but null/undersized"}, + }; + for (const auto& test_case : cases) { + SCOPED_TRACE(test_case.name); + expect_build_error(test_case.spec, test_case.error); + } +} + +// A live shape whose element count is not a multiple of K must throw instead of +// dispatching a mis-sized M=3 grid, and the graph must stay retryable. +TEST(Q4gswM3FailsClosed, RejectsALiveShapeThatBreaksTheKContract) { + Q4gswSpec spec; + spec.input_dims = {static_cast(kMaxM), static_cast(kK)}; + spec.dynamic = true; + WebGPUGraph graph; + Q4gswIds ids; + ASSERT_NO_THROW(ids = build_q4gsw_graph(graph, spec)); + + graph.resize_input(ids.input, {3, 8}); + std::string error; + try { + graph.propagate_resize(); + } catch (const std::exception& exception) { + error = exception.what(); + } + EXPECT_EQ( + error, + "WebGPU linear_q4gsw(resize): live input numel not a multiple of K"); + + graph.resize_input(ids.input, {3, kK}); + ASSERT_NO_THROW(graph.propagate_resize()); + const ActiveDispatch active = active_dispatch(graph); + EXPECT_EQ(active.count, 1u); + EXPECT_EQ(active.kernel_name, kM3Kernel); +} + +} // namespace executorch::backends::webgpu + +int main(int argc, char** argv) { + ::testing::InitGoogleTest(&argc, argv); + + executorch::backends::webgpu::WebGPUContext ctx; + try { + ctx = executorch::backends::webgpu::create_webgpu_context(); + } catch (const std::exception& e) { + if (std::getenv("WEBGPU_REQUIRE_DEVICE") != nullptr) { + std::printf( + "FAIL: WEBGPU_REQUIRE_DEVICE set but no device: %s\n", e.what()); + return 1; + } + std::printf("SKIP: %s\n", e.what()); + return 0; + } + executorch::backends::webgpu::set_default_webgpu_context(&ctx); + executorch::backends::webgpu::g_device = ctx.device; + std::printf("WebGPU device acquired (native)\n"); + + const int rc = RUN_ALL_TESTS(); + executorch::backends::webgpu::set_default_webgpu_context(nullptr); + executorch::backends::webgpu::destroy_webgpu_context(ctx); + return rc; +} diff --git a/backends/webgpu/test/native/test_scatter.cpp b/backends/webgpu/test/native/test_scatter.cpp new file mode 100644 index 00000000000..8473f8f62c3 --- /dev/null +++ b/backends/webgpu/test/native/test_scatter.cpp @@ -0,0 +1,401 @@ +/* + * 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 +#include +#include +#include +#include +#include +#include +#include +#include + +namespace executorch::backends::webgpu { +namespace { + +constexpr int64_t kVocabSize = 262144; +constexpr int64_t kSelectedCount = 4096; +constexpr uint32_t kUniqueLanes = 64; + +constexpr const char* kGenericOp = "aten.scatter.src"; +constexpr const char* kUniqueOp = "et_vk.scatter_src_unique.default"; + +struct FixtureCase { + std::string name; + bool parallel_equivalent = false; + bool official_provenance = false; +}; + +// NOLINTBEGIN(cppcoreguidelines-avoid-non-const-global-variables) +std::string g_dir; +WGPUDevice g_device = nullptr; +// NOLINTEND(cppcoreguidelines-avoid-non-const-global-variables) + +template +std::vector read_bin(const std::string& path, size_t count) { + std::ifstream f(path, std::ios::binary | std::ios::ate); + if (!f || static_cast(f.tellg()) != count * sizeof(T)) { + return {}; + } + f.seekg(0); + std::vector data(count); + f.read( + reinterpret_cast(data.data()), + static_cast(count * sizeof(T))); + return data; +} + +std::vector read_cases() { + std::vector cases; + std::ifstream manifest(g_dir + "/cases.txt"); + std::string line; + while (std::getline(manifest, line)) { + if (line.empty()) { + continue; + } + std::istringstream parts(line); + FixtureCase entry; + int equivalent = 0; + int provenance = 0; + parts >> entry.name >> equivalent >> provenance; + entry.parallel_equivalent = equivalent != 0; + entry.official_provenance = provenance != 0; + cases.push_back(entry); + } + return cases; +} + +struct ScatterGraphSpec { + const char* op = kGenericOp; + int64_t vocab = kVocabSize; + int64_t selected = kSelectedCount; + int64_t output_vocab = kVocabSize; + int64_t dim = -1; + bool index_is_int = true; + bool source_is_int = false; + bool input_is_int = false; + bool drop_last_arg = false; +}; + +void build_scatter_graph(WebGPUGraph& graph, const ScatterGraphSpec& spec) { + namespace vk = vkgraph; + ::flatbuffers::FlatBufferBuilder fbb; + std::vector<::flatbuffers::Offset> values; + auto add_tensor = [&](vk::VkDataType dtype, int64_t width, int mem_obj_id) { + const int id = static_cast(values.size()); + const std::vector dims = {1, 1, static_cast(width)}; + values.push_back(vk::CreateVkValue( + fbb, + vk::GraphTypes::VkTensor, + vk::CreateVkTensorDirect( + fbb, dtype, &dims, /*constant_id=*/-1, mem_obj_id) + .Union())); + return id; + }; + + const int input = add_tensor( + spec.input_is_int ? vk::VkDataType::INT32 : vk::VkDataType::FLOAT32, + spec.vocab, + 0); + const int dim = static_cast(values.size()); + values.push_back(vk::CreateVkValue( + fbb, vk::GraphTypes::Int, vk::CreateInt(fbb, spec.dim).Union())); + const int index = add_tensor( + spec.index_is_int ? vk::VkDataType::INT32 : vk::VkDataType::FLOAT32, + spec.selected, + 1); + const int source = add_tensor( + spec.source_is_int ? vk::VkDataType::INT32 : vk::VkDataType::FLOAT32, + spec.selected, + 2); + const int output = add_tensor(vk::VkDataType::FLOAT32, spec.output_vocab, 3); + + std::vector args = {input, dim, index, source, output}; + if (spec.drop_last_arg) { + args.pop_back(); + } + std::vector<::flatbuffers::Offset> chain; + chain.push_back(vk::CreateOperatorCallDirect(fbb, 0, spec.op, &args)); + + const std::vector input_ids = { + static_cast(input), + static_cast(index), + static_cast(source)}; + const std::vector output_ids = {static_cast(output)}; + const auto root = vk::CreateVkGraphDirect( + fbb, "0", &chain, &values, &input_ids, &output_ids); + vk::FinishVkGraphBuffer(fbb, root); + + graph.set_device(g_device); + graph.build(fbb.GetBufferPointer(), nullptr, 0, nullptr); +} + +// The compute dispatch, skipping the flat input->output copy the handler emits. +const WebGPUDispatch& compute_dispatch(WebGPUGraph& graph) { + for (size_t i = 0; i < graph.num_dispatches(); i++) { + if (graph.dispatch_at(i).kind == WebGPUDispatch::Kind::Compute) { + return graph.dispatch_at(i); + } + } + throw std::runtime_error("scatter graph has no compute dispatch"); +} + +std::vector run_route( + const char* op, + const std::vector& base, + const std::vector& indices, + const std::vector& source) { + ScatterGraphSpec spec; + spec.op = op; + WebGPUGraph graph; + build_scatter_graph(graph, spec); + + std::vector inputs(3); + inputs[0] = {base.data(), base.size() * sizeof(float), false, true}; + inputs[1] = {indices.data(), indices.size() * sizeof(int32_t), false, false}; + inputs[2] = {source.data(), source.size() * sizeof(float), false, true}; + std::vector out(static_cast(kVocabSize), 0.0f); + std::vector outputs(1); + outputs[0] = {out.data(), out.size() * sizeof(float), /*fp32=*/true}; + + graph.copy_inputs(inputs); + const WebGPUExecutionPlan plan = graph.make_execution_plan({}); + graph.execute(plan); + graph.copy_outputs(outputs, plan); + return out; +} + +} // namespace + +TEST(ScatterFixtureContract, ExportedCorpusIsPresentAndLabelled) { + const std::vector cases = read_cases(); + ASSERT_FALSE(cases.empty()) << "missing " << g_dir << "/cases.txt"; + size_t provenance_count = 0; + size_t equivalent_count = 0; + for (const FixtureCase& entry : cases) { + equivalent_count += entry.parallel_equivalent ? 1 : 0; + provenance_count += entry.official_provenance ? 1 : 0; + EXPECT_FALSE(entry.official_provenance && !entry.parallel_equivalent) + << entry.name << ": provenance must imply parallel equivalence"; + } + EXPECT_GT(provenance_count, 0u); + EXPECT_GT(equivalent_count, provenance_count) + << "corpus must contain an equivalent-but-uncertifiable case"; + EXPECT_LT(provenance_count, cases.size()); +} + +TEST(ScatterRoute, GenericAndUniqueAreDistinguishableWitnesses) { + WebGPUGraph generic; + ScatterGraphSpec generic_spec; + build_scatter_graph(generic, generic_spec); + const WebGPUDispatch& generic_dispatch = compute_dispatch(generic); + EXPECT_EQ(generic_dispatch.kernel_name, "scatter"); + EXPECT_EQ(generic_dispatch.workgroup_count_x, 1u); + + WebGPUGraph unique; + ScatterGraphSpec unique_spec; + unique_spec.op = kUniqueOp; + build_scatter_graph(unique, unique_spec); + const WebGPUDispatch& unique_dispatch = compute_dispatch(unique); + EXPECT_EQ(unique_dispatch.kernel_name, "scatter_unique_indices"); + EXPECT_EQ( + unique_dispatch.workgroup_count_x, + static_cast(kSelectedCount) / kUniqueLanes); + + // The forced-WG1 control: a witness that reports the parallel route cannot be + // satisfied by the serial one. + EXPECT_NE(generic_dispatch.kernel_name, unique_dispatch.kernel_name); + EXPECT_NE( + generic_dispatch.workgroup_count_x, unique_dispatch.workgroup_count_x); + EXPECT_EQ( + get_webgpu_shader_info("scatter_unique_indices").workgroup_size_x, + kUniqueLanes); + EXPECT_EQ(get_webgpu_shader_info("scatter").workgroup_size_x, 1u); +} + +TEST(ScatterExactness, MatchesTheCpuAuthorityOnEveryCase) { + const std::vector base = + read_bin(g_dir + "/base.bin", kVocabSize); + ASSERT_FALSE(base.empty()) << "missing " << g_dir << "/base.bin"; + + for (const FixtureCase& entry : read_cases()) { + SCOPED_TRACE(entry.name); + const std::string prefix = g_dir + "/" + entry.name; + const std::vector indices = + read_bin(prefix + ".index.bin", kSelectedCount); + const std::vector source = + read_bin(prefix + ".source.bin", kSelectedCount); + const std::vector expected = + read_bin(prefix + ".expected.bin", kVocabSize); + ASSERT_FALSE(indices.empty()); + ASSERT_FALSE(source.empty()); + ASSERT_FALSE(expected.empty()); + + const std::vector generic = + run_route(kGenericOp, base, indices, source); + ASSERT_EQ(generic.size(), expected.size()); + for (size_t i = 0; i < expected.size(); i++) { + ASSERT_EQ(generic[i], expected[i]) << "generic mismatch at " << i; + } + + if (!entry.parallel_equivalent) { + continue; + } + const std::vector unique = + run_route(kUniqueOp, base, indices, source); + ASSERT_EQ(unique.size(), expected.size()); + for (size_t i = 0; i < expected.size(); i++) { + ASSERT_EQ(unique[i], expected[i]) << "unique mismatch at " << i; + } + } +} + +TEST(ScatterExactness, DuplicateDestinationsStayOnTheDuplicateSafeRoute) { + const std::vector base = + read_bin(g_dir + "/base.bin", kVocabSize); + ASSERT_FALSE(base.empty()); + const std::string prefix = g_dir + "/duplicate_destinations"; + const std::vector indices = + read_bin(prefix + ".index.bin", kSelectedCount); + const std::vector source = + read_bin(prefix + ".source.bin", kSelectedCount); + const std::vector expected = + read_bin(prefix + ".expected.bin", kVocabSize); + ASSERT_FALSE(indices.empty() || source.empty() || expected.empty()); + + const std::vector generic = + run_route(kGenericOp, base, indices, source); + for (size_t i = 0; i < expected.size(); i++) { + ASSERT_EQ(generic[i], expected[i]) << "last-write-wins broken at " << i; + } +} + +TEST(ScatterFailsClosed, RejectsMalformedShapesAndScalars) { + for (const char* op : {kGenericOp, kUniqueOp}) { + const struct { + const char* name; + ScatterGraphSpec spec; + } cases[] = { + {"dim != -1", + [op] { + ScatterGraphSpec s; + s.op = op; + s.dim = 0; + return s; + }()}, + {"input width", + [op] { + ScatterGraphSpec s; + s.op = op; + s.vocab = 1024; + return s; + }()}, + {"selected width", + [op] { + ScatterGraphSpec s; + s.op = op; + s.selected = 2048; + return s; + }()}, + {"output width", + [op] { + ScatterGraphSpec s; + s.op = op; + s.output_vocab = kVocabSize / 2; + return s; + }()}, + {"index dtype", + [op] { + ScatterGraphSpec s; + s.op = op; + s.index_is_int = false; + return s; + }()}, + {"source dtype", + [op] { + ScatterGraphSpec s; + s.op = op; + s.source_is_int = true; + return s; + }()}, + {"input dtype", + [op] { + ScatterGraphSpec s; + s.op = op; + s.input_is_int = true; + return s; + }()}, + {"argument count", + [op] { + ScatterGraphSpec s; + s.op = op; + s.drop_last_arg = true; + return s; + }()}, + }; + for (const auto& test_case : cases) { + SCOPED_TRACE(std::string(op) + " / " + test_case.name); + WebGPUGraph graph; + EXPECT_THROW( + build_scatter_graph(graph, test_case.spec), std::runtime_error); + } + } +} + +TEST(ScatterFailsClosed, AcceptsTheExactMtpShapeOnBothRoutes) { + for (const char* op : {kGenericOp, kUniqueOp}) { + SCOPED_TRACE(op); + ScatterGraphSpec spec; + spec.op = op; + WebGPUGraph graph; + EXPECT_NO_THROW(build_scatter_graph(graph, spec)); + } +} + +} // namespace executorch::backends::webgpu + +int main(int argc, char** argv) { + ::testing::InitGoogleTest(&argc, argv); + + executorch::backends::webgpu::g_dir = "/tmp/scatter"; + if (argc > 1) { + executorch::backends::webgpu::g_dir = argv[1]; + } + if (const char* env = std::getenv("WEBGPU_SCATTER_DIR")) { + executorch::backends::webgpu::g_dir = env; + } + + executorch::backends::webgpu::WebGPUContext ctx; + try { + ctx = executorch::backends::webgpu::create_webgpu_context(); + } catch (const std::exception& e) { + if (std::getenv("WEBGPU_REQUIRE_DEVICE") != nullptr) { + std::printf( + "FAIL: WEBGPU_REQUIRE_DEVICE set but no device: %s\n", e.what()); + return 1; + } + std::printf("SKIP: %s\n", e.what()); + return 0; + } + executorch::backends::webgpu::set_default_webgpu_context(&ctx); + executorch::backends::webgpu::g_device = ctx.device; + std::printf("WebGPU device acquired (native)\n"); + + const int rc = RUN_ALL_TESTS(); + executorch::backends::webgpu::set_default_webgpu_context(nullptr); + executorch::backends::webgpu::destroy_webgpu_context(ctx); + return rc; +} diff --git a/backends/webgpu/test/native/test_topk.cpp b/backends/webgpu/test/native/test_topk.cpp new file mode 100644 index 00000000000..460ec14214c --- /dev/null +++ b/backends/webgpu/test/native/test_topk.cpp @@ -0,0 +1,402 @@ +/* + * 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 +#include +#include +#include +#include +#include +#include +#include +#include + +namespace executorch::backends::webgpu { +namespace { + +constexpr int64_t kInputWidth = 2048; +constexpr int64_t kOutputWidth = 32; +constexpr uint32_t kTopkStagedLanes = 64; + +// Must equal the corpus written by test/ops/topk/export_topk_artifacts.py. +const char* const kCases[] = { + "all_equal", + // all_negative and straddling_zero are the only rows reaching topk.wgsl:62. + "all_negative", + "boundary_ties", + "descending", + "infinities", + "interior_ties", + "nan_payloads", + "ordinary", + "random_seeded", + "signed_zeros", + "straddling_zero", +}; + +// NOLINTBEGIN(cppcoreguidelines-avoid-non-const-global-variables) +std::string g_dir; +WGPUDevice g_device = nullptr; +// NOLINTEND(cppcoreguidelines-avoid-non-const-global-variables) + +std::vector read_words(const std::string& path, size_t count) { + std::ifstream f(path, std::ios::binary | std::ios::ate); + if (!f || static_cast(f.tellg()) != count * sizeof(uint32_t)) { + return {}; + } + f.seekg(0); + std::vector data(count); + f.read( + reinterpret_cast(data.data()), + static_cast(count * sizeof(uint32_t))); + return data; +} + +// Strips `//` and `/* */` so a shader's prose cannot shadow its real attribute. +std::string strip_wgsl_comments(const std::string& src) { + std::string out; + out.reserve(src.size()); + for (size_t i = 0; i < src.size();) { + if (src.compare(i, 2, "//") == 0) { + while (i < src.size() && src[i] != '\n') { + i++; + } + } else if (src.compare(i, 2, "/*") == 0) { + i += 2; + while (i + 1 < src.size() && src.compare(i, 2, "*/") != 0) { + i++; + } + i = i + 1 < src.size() ? i + 2 : src.size(); + } else { + out.push_back(src[i++]); + } + } + return out; +} + +// The declared x dim, resolving a `const : u32 = u;` indirection. +uint32_t declared_workgroup_size_x(const std::string& src) { + const std::string code = strip_wgsl_comments(src); + const size_t at = code.find("@workgroup_size"); + if (at == std::string::npos) { + return 0; + } + const size_t open = code.find('(', at); + const size_t close = code.find(')', open); + if (open == std::string::npos || close == std::string::npos) { + return 0; + } + std::string token = code.substr(open + 1, close - open - 1); + token = token.substr(0, token.find(',')); + const size_t begin = token.find_first_not_of(" \t\n\r"); + if (begin == std::string::npos) { + return 0; + } + token = token.substr(begin, token.find_last_not_of(" \t\n\r") + 1 - begin); + if (!token.empty() && std::isdigit(static_cast(token[0]))) { + return static_cast(std::strtoul(token.c_str(), nullptr, 10)); + } + const size_t decl = code.find("const " + token); + const size_t eq = + decl == std::string::npos ? std::string::npos : code.find('=', decl); + if (eq == std::string::npos) { + return 0; + } + return static_cast( + std::strtoul(code.c_str() + eq + 1, nullptr, 10)); +} + +struct TopkGraphSpec { + int64_t input_width = kInputWidth; + int64_t values_width = kOutputWidth; + int64_t indices_width = kOutputWidth; + int64_t k = kOutputWidth; + int64_t dim = -1; + bool largest = true; + bool sorted = true; + bool values_are_int = false; + bool indices_are_int = true; + bool output_list_is_int = false; + bool drop_last_arg = false; +}; + +void build_topk_graph(WebGPUGraph& graph, const TopkGraphSpec& spec) { + namespace vk = vkgraph; + ::flatbuffers::FlatBufferBuilder fbb; + std::vector<::flatbuffers::Offset> values; + auto add_tensor = [&](vk::VkDataType dtype, + const std::vector& dims, + int mem_obj_id) { + const int id = static_cast(values.size()); + values.push_back(vk::CreateVkValue( + fbb, + vk::GraphTypes::VkTensor, + vk::CreateVkTensorDirect( + fbb, dtype, &dims, /*constant_id=*/-1, mem_obj_id) + .Union())); + return id; + }; + auto add_int = [&](int64_t value) { + const int id = static_cast(values.size()); + values.push_back(vk::CreateVkValue( + fbb, vk::GraphTypes::Int, vk::CreateInt(fbb, value).Union())); + return id; + }; + auto add_bool = [&](bool value) { + const int id = static_cast(values.size()); + values.push_back(vk::CreateVkValue( + fbb, vk::GraphTypes::Bool, vk::CreateBool(fbb, value).Union())); + return id; + }; + + const int input = add_tensor( + vk::VkDataType::FLOAT32, + {1, 1, static_cast(spec.input_width)}, + 0); + const int k = add_int(spec.k); + const int dim = add_int(spec.dim); + const int largest = add_bool(spec.largest); + const int sorted = add_bool(spec.sorted); + const int out_values = add_tensor( + spec.values_are_int ? vk::VkDataType::INT32 : vk::VkDataType::FLOAT32, + {1, 1, static_cast(spec.values_width)}, + 1); + const int out_indices = add_tensor( + spec.indices_are_int ? vk::VkDataType::INT32 : vk::VkDataType::FLOAT32, + {1, 1, static_cast(spec.indices_width)}, + 2); + + int output_list = 0; + if (spec.output_list_is_int) { + output_list = add_int(0); + } else { + const std::vector items = {out_values, out_indices}; + output_list = static_cast(values.size()); + values.push_back(vk::CreateVkValue( + fbb, + vk::GraphTypes::ValueList, + vk::CreateValueListDirect(fbb, &items).Union())); + } + + std::vector args = {input, k, dim, largest, sorted, output_list}; + if (spec.drop_last_arg) { + args.pop_back(); + } + std::vector<::flatbuffers::Offset> chain; + chain.push_back( + vk::CreateOperatorCallDirect(fbb, 0, "aten.topk.default", &args)); + + const std::vector input_ids = {static_cast(input)}; + const std::vector output_ids = { + static_cast(out_values), static_cast(out_indices)}; + const auto root = vk::CreateVkGraphDirect( + fbb, "0", &chain, &values, &input_ids, &output_ids); + vk::FinishVkGraphBuffer(fbb, root); + + graph.set_device(g_device); + graph.build(fbb.GetBufferPointer(), nullptr, 0, nullptr); +} + +void run_case(const char* name) { + const std::string base = g_dir + "/" + name; + const std::vector scores = + read_words(base + ".scores.bin", kInputWidth); + const std::vector expected_values = + read_words(base + ".values.bin", kOutputWidth); + const std::vector expected_indices = + read_words(base + ".indices.bin", kOutputWidth); + ASSERT_FALSE(scores.empty()) << "missing/short scores fixture for " << name; + ASSERT_FALSE(expected_values.empty()) << "missing values fixture " << name; + ASSERT_FALSE(expected_indices.empty()) << "missing indices fixture " << name; + + WebGPUGraph graph; + build_topk_graph(graph, TopkGraphSpec{}); + ASSERT_EQ(graph.num_dispatches(), 1u); + EXPECT_EQ(graph.dispatch_at(0).kernel_name, "topk_staged_serial"); + EXPECT_EQ(graph.dispatch_at(0).workgroup_count_x, 1u); + + std::vector inputs(1); + inputs[0] = {scores.data(), scores.size() * sizeof(uint32_t), false, true}; + std::vector got_values(kOutputWidth, 0u); + std::vector got_indices(kOutputWidth, -1); + std::vector outputs(2); + outputs[0] = { + got_values.data(), got_values.size() * sizeof(uint32_t), /*fp32=*/true}; + outputs[1] = { + got_indices.data(), got_indices.size() * sizeof(int32_t), /*fp32=*/false}; + + graph.copy_inputs(inputs); + const WebGPUExecutionPlan plan = graph.make_execution_plan({}); + graph.execute(plan); + graph.copy_outputs(outputs, plan); + + for (int64_t i = 0; i < kOutputWidth; i++) { + // Bit-exact: NaN payloads and signed zeros must survive unchanged. + EXPECT_EQ(got_values[i], expected_values[i]) << name << " value " << i; + EXPECT_EQ(static_cast(got_indices[i]), expected_indices[i]) + << name << " index " << i; + } +} + +} // namespace + +TEST(TopkFixtureContract, CaseListMatchesTheExportedCorpus) { + std::ifstream manifest(g_dir + "/cases.txt"); + ASSERT_TRUE(manifest.good()) << "missing " << g_dir << "/cases.txt"; + std::vector exported; + std::string line; + while (std::getline(manifest, line)) { + if (!line.empty()) { + exported.push_back(line); + } + } + std::vector expected(std::begin(kCases), std::end(kCases)); + EXPECT_EQ(exported, expected); +} + +TEST(TopkShader, RegistryWorkgroupSizeMatchesTheShaderDeclaration) { + const WebGPUShaderInfo& info = get_webgpu_shader_info("topk"); + ASSERT_NE(info.source, nullptr); + EXPECT_EQ(declared_workgroup_size_x(info.source), kTopkStagedLanes); + EXPECT_EQ(info.workgroup_size_x, kTopkStagedLanes) + << "registry workgroup size disagrees with @workgroup_size; the " + "generated header parsed a commented-out attribute"; +} + +TEST(TopkExactness, AllCases) { + for (const char* name : kCases) { + SCOPED_TRACE(name); + run_case(name); + } +} + +TEST(TopkFailsClosed, RejectsMalformedShapesAndScalars) { + const struct { + const char* name; + TopkGraphSpec spec; + } cases[] = { + {"k != 32", + [] { + TopkGraphSpec s; + s.k = 16; + return s; + }()}, + {"dim != -1", + [] { + TopkGraphSpec s; + s.dim = 2; + return s; + }()}, + {"largest = false", + [] { + TopkGraphSpec s; + s.largest = false; + return s; + }()}, + {"sorted = false", + [] { + TopkGraphSpec s; + s.sorted = false; + return s; + }()}, + {"input width", + [] { + TopkGraphSpec s; + s.input_width = 1024; + return s; + }()}, + {"values width", + [] { + TopkGraphSpec s; + s.values_width = 64; + return s; + }()}, + {"indices width", + [] { + TopkGraphSpec s; + s.indices_width = 31; + return s; + }()}, + {"values dtype", + [] { + TopkGraphSpec s; + s.values_are_int = true; + return s; + }()}, + {"indices dtype", + [] { + TopkGraphSpec s; + s.indices_are_int = false; + return s; + }()}, + {"output list type", + [] { + TopkGraphSpec s; + s.output_list_is_int = true; + return s; + }()}, + {"argument count", + [] { + TopkGraphSpec s; + s.drop_last_arg = true; + return s; + }()}, + }; + for (const auto& test_case : cases) { + SCOPED_TRACE(test_case.name); + WebGPUGraph graph; + EXPECT_THROW(build_topk_graph(graph, test_case.spec), std::runtime_error); + } +} + +TEST(TopkFailsClosed, AcceptsTheExactMtpShape) { + WebGPUGraph graph; + EXPECT_NO_THROW(build_topk_graph(graph, TopkGraphSpec{})); +} + +} // namespace executorch::backends::webgpu + +int main(int argc, char** argv) { + ::testing::InitGoogleTest(&argc, argv); + + executorch::backends::webgpu::g_dir = "/tmp/topk"; + if (argc > 1) { + executorch::backends::webgpu::g_dir = argv[1]; + } + if (const char* env = std::getenv("WEBGPU_TOPK_DIR")) { + executorch::backends::webgpu::g_dir = env; + } + + executorch::backends::webgpu::WebGPUContext ctx; + try { + ctx = executorch::backends::webgpu::create_webgpu_context(); + } catch (const std::exception& e) { + if (std::getenv("WEBGPU_REQUIRE_DEVICE") != nullptr) { + std::printf( + "FAIL: WEBGPU_REQUIRE_DEVICE set but no device: %s\n", e.what()); + return 1; + } + std::printf("SKIP: %s\n", e.what()); + return 0; + } + executorch::backends::webgpu::set_default_webgpu_context(&ctx); + executorch::backends::webgpu::g_device = ctx.device; + std::printf("WebGPU device acquired (native)\n"); + + const int rc = RUN_ALL_TESTS(); + executorch::backends::webgpu::set_default_webgpu_context(nullptr); + executorch::backends::webgpu::destroy_webgpu_context(ctx); + return rc; +} diff --git a/backends/webgpu/test/op_tests/cases.py b/backends/webgpu/test/op_tests/cases.py index a96a96ab56f..a146ac722e2 100644 --- a/backends/webgpu/test/op_tests/cases.py +++ b/backends/webgpu/test/op_tests/cases.py @@ -48,9 +48,14 @@ CONFIGS as _CAT_CONFIGS, ) from executorch.backends.webgpu.test.ops.test_compare import ( + _det_input as _scalar_compare_input, compare_gen_a, compare_gen_b, CompareModule, + SCALAR, + SCALAR_OPS, + SCALAR_SHAPES, + ScalarCompareModule, ) from executorch.backends.webgpu.test.ops.test_conv1d_dw import Conv1dDWModule from executorch.backends.webgpu.test.ops.test_conv1d_pw import ( @@ -168,6 +173,7 @@ to_copy_float_input, to_copy_int_input, ToCopyBoolToFloatModule, + ToCopyFloatToInt64Module, ToCopyFloatToIntToFloatModule, ToCopyIntToFloatModule, ) @@ -1773,6 +1779,9 @@ def _relu_suite() -> WebGPUTestSuite: from executorch.backends.webgpu.test.ops.test_sub import ( CONFIGS as _SUB_CONFIGS, + int32_wrap_input_a, + int32_wrap_input_b, + SubAlphaModule, SubModule, ) @@ -1783,9 +1792,38 @@ def _sub_suite() -> WebGPUTestSuite: # over a TensorMeta UBO); fp64 golden. Mirrors _mul_suite. alpha is a # construct kwarg baked into the .pte, never a serialized input. return WebGPUTestSuite( - module_factory=lambda: SubModule(), + module_factory=lambda alpha=None: ( + SubModule() if alpha is None else SubAlphaModule(alpha) + ), cases=[ - Case(name=name, inputs=(sa, sb)) for name, (sa, sb) in _SUB_CONFIGS.items() + *[ + Case(name=name, inputs=(sa, sb)) + for name, (sa, sb) in _SUB_CONFIGS.items() + ], + Case( + name="int32_wrap", + construct={"alpha": 1}, + inputs=( + InputSpec(shape=(2, 4), gen=int32_wrap_input_a), + InputSpec(shape=(2, 4), gen=int32_wrap_input_b), + ), + ), + Case( + name="int32_broadcast", + construct={"alpha": 1}, + inputs=( + InputSpec(shape=(2, 3), gen=int32_wrap_input_a), + InputSpec(shape=(1, 3), gen=int32_wrap_input_b), + ), + ), + Case( + name="int32_alpha", + construct={"alpha": 3}, + inputs=( + InputSpec(shape=(2, 4), gen=int32_wrap_input_a), + InputSpec(shape=(2, 4), gen=int32_wrap_input_b), + ), + ), ], ) @@ -2218,3 +2256,145 @@ def _argmin_suite() -> WebGPUTestSuite: ], golden_dtype="float32", ) + + +@register_op_test("to_copy_f2i") +def _to_copy_f2i_suite() -> WebGPUTestSuite: + # Terminal float->int cast; i32() truncates toward zero, so the .75 magnitudes bite. + return WebGPUTestSuite( + module_factory=ToCopyFloatToInt64Module, + cases=[ + Case( + name=f"length_{n}", + inputs=(InputSpec(shape=(n,), gen=to_copy_float_input),), + ) + for n in (63, 64, 65, 257) + ], + golden_dtype="float32", + ) + + +from executorch.backends.webgpu.test.ops.test_gather import ( + CONFIGS as _GATHER_CONFIGS, + gather_index_gen, + gather_self_gen, + GatherModule, +) + + +@register_op_test("gather") +def _gather_suite() -> WebGPUTestSuite: + # out/index share a shape; rank3_neg's dim=-1 pins the handler's dim normalization. + return WebGPUTestSuite( + module_factory=lambda dim: GatherModule(dim), + cases=[ + Case( + name=name, + construct={"dim": dim}, + inputs=( + InputSpec(shape=self_shape, gen=gather_self_gen), + InputSpec(shape=index_shape, gen=gather_index_gen(self_shape[dim])), + ), + ) + for name, (self_shape, dim, index_shape) in _GATHER_CONFIGS.items() + ], + golden_dtype="float32", + ) + + +from executorch.backends.webgpu.test.ops.test_where import ( + CONFIGS as _WHERE_CONFIGS, + where_a_gen, + where_b_gen, + where_cond_gen, + WhereModule, +) + + +@register_op_test("where") +def _where_suite() -> WebGPUTestSuite: + # bool cond + fp32 a/b, all broadcast; a>0 and b<0 so a wrong pick flips sign. + return WebGPUTestSuite( + module_factory=lambda: WhereModule(), + cases=[ + Case( + name=name, + inputs=( + InputSpec(shape=cond_shape, gen=where_cond_gen), + InputSpec(shape=a_shape, gen=where_a_gen), + InputSpec(shape=b_shape, gen=where_b_gen), + ), + ) + for name, (cond_shape, a_shape, b_shape) in _WHERE_CONFIGS.items() + ], + golden_dtype="float32", + ) + + +@register_op_test("compare_scalar") +def _compare_scalar_suite() -> WebGPUTestSuite: + # All six .Scalar variants vs 0.0; tail numel 15 hits the partial bool word. + return WebGPUTestSuite( + module_factory=lambda op, scalar: ScalarCompareModule(op, scalar), + cases=[ + Case( + name=f"{op}_{shape_name}", + construct={"op": op, "scalar": SCALAR}, + inputs=(InputSpec(shape=shape, gen=_scalar_compare_input),), + ) + for op in SCALAR_OPS + for shape_name, shape in SCALAR_SHAPES.items() + ], + golden_dtype="bool", + ) + + +from executorch.backends.webgpu.test.ops.test_logical_not import ( + _det_input as _logical_not_input, + CONFIGS as _LOGICAL_NOT_CONFIGS, + LogicalNotModule, +) + + +@register_op_test("logical_not") +def _logical_not_suite() -> WebGPUTestSuite: + # not(x >= 0) over a GPU-derived mask; tail3x7 numel 21 hits the partial bool word. + return WebGPUTestSuite( + module_factory=lambda threshold: LogicalNotModule(threshold), + cases=[ + Case( + name=name, + construct={"threshold": 0.0}, + inputs=(InputSpec(shape=shape, gen=_logical_not_input),), + ) + for name, shape in _LOGICAL_NOT_CONFIGS.items() + ], + golden_dtype="bool", + ) + + +from executorch.backends.webgpu.test.ops.index.test_index import ( + CONFIGS as _INDEX_CONFIGS, + index_idx_gen, + index_self_gen, + IndexModule, +) + + +@register_op_test("index") +def _index_suite() -> WebGPUTestSuite: + # 1D-self index out[i]=self[idx[i]]; repeat/reverse idx exposes a bad gather. + return WebGPUTestSuite( + module_factory=lambda: IndexModule(), + cases=[ + Case( + name=name, + inputs=( + InputSpec(shape=(self_len,), gen=index_self_gen), + InputSpec(shape=(len(idx),), gen=index_idx_gen(idx)), + ), + ) + for name, (self_len, idx) in _INDEX_CONFIGS.items() + ], + golden_dtype="float32", + ) diff --git a/backends/webgpu/test/op_tests/driver_util.h b/backends/webgpu/test/op_tests/driver_util.h index 96998cec49e..2c24304c6f2 100644 --- a/backends/webgpu/test/op_tests/driver_util.h +++ b/backends/webgpu/test/op_tests/driver_util.h @@ -24,7 +24,7 @@ struct GoldenRef { std::string path; std::vector shape; int output_index = 0; - std::string dtype = "float32"; // "float32" | "int8" | "int64" (argmax index) + std::string dtype = "float32"; // "float32" | "int8" | "int32" | "int64" }; struct ManifestEntry { diff --git a/backends/webgpu/test/op_tests/generate_op_tests.py b/backends/webgpu/test/op_tests/generate_op_tests.py index 33b7940a9be..600e2ed0e12 100644 --- a/backends/webgpu/test/op_tests/generate_op_tests.py +++ b/backends/webgpu/test/op_tests/generate_op_tests.py @@ -92,6 +92,10 @@ def _write_int64(t: torch.Tensor, path: str) -> None: t.detach().contiguous().cpu().numpy().astype(" None: + t.detach().contiguous().cpu().numpy().astype("(); + int mism = -1; + for (size_t i = 0; i < gn; i++) { + if (out_p[i] != golden[i]) { + mism = static_cast(i); + break; + } + } + EXPECT_EQ(mism, -1) << "int32 mismatch at index " << mism + << ": out=" << (mism >= 0 ? out_p[mism] : 0) + << " golden=" << (mism >= 0 ? golden[mism] : 0); } else if (e_.golden.dtype == "int64") { // int64-index ops (argmax/argmin) compare exact: indices are discrete. auto golden = load_int64_bin(e_.golden.path, gn); diff --git a/backends/webgpu/test/op_tests/test_typed_input_contract.py b/backends/webgpu/test/op_tests/test_typed_input_contract.py new file mode 100644 index 00000000000..2624d308685 --- /dev/null +++ b/backends/webgpu/test/op_tests/test_typed_input_contract.py @@ -0,0 +1,179 @@ +# 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. + +# pyre-strict + +"""Source contract for the op-test driver's typed input loader. + +`op_test_driver.cpp` loads each manifest input by its declared dtype -- `bool` via +`load_int8_bin` + `ScalarType::Bool`, `int32` via `load_int32_bin`, everything else +via `load_fp32_bin`. Deleting the BOOL branch, reordering it after the fp32 +fallback, or widening a bool input to fp32 would silently change what every +bool-input case (`where`, `to_copy_bool_input_to_float`) actually feeds the GPU, +and the golden comparison would still "pass" on the wrong bytes. + +These checks read source, not artifacts: they need no torch and no GPU, so they run +in any environment. The producer side (`generate_op_tests.py`) is pinned to the same +three dtype strings, and `cases.py` is pinned to wire at least one input generator +per dtype -- i.e. a generated manifest carries at least one entry of each. +""" + +from __future__ import annotations + +import ast +import re +import unittest +from pathlib import Path + +_OP_TESTS_DIR: Path = Path(__file__).resolve().parent +_OPS_DIR: Path = _OP_TESTS_DIR.parent / "ops" + +_DRIVER: Path = _OP_TESTS_DIR / "op_test_driver.cpp" +_GENERATOR: Path = _OP_TESTS_DIR / "generate_op_tests.py" +_CASES: Path = _OP_TESTS_DIR / "cases.py" + +_INPUT_LOOP_START: str = "for (const auto& in : e_.inputs) {" +_INPUT_LOOP_END: str = "std::vector inputs;" +_BOOL_BRANCH: str = 'if (in.dtype == "bool") {' +_INT32_BRANCH: str = '} else if (in.dtype == "int32") {' +_FP32_BRANCH: str = "} else {" + +_WRITE_LOOP_START: str = "for i, t in enumerate(inputs):" +_WRITE_LOOP_END: str = "input_entries.append(" + +# manifest input dtype -> (cases.py generator, defining module, dtype literal) +_DTYPE_GENERATORS: dict[str, tuple[str, Path, str]] = { + "bool": ("where_cond_gen", _OPS_DIR / "test_where.py", "torch.bool"), + "int32": ("to_copy_int_input", _OPS_DIR / "test_to_copy.py", "torch.int32"), + "float32": ("to_copy_float_input", _OPS_DIR / "test_to_copy.py", "torch.float32"), +} + +# The op names the D10 artifact command passes to `generate_op_tests --ops`. +_D10_SUITES: tuple[str, ...] = ( + "split_with_sizes_copy", + "to_copy", + "to_copy_f2i", + "expand_copy", + "gather", + "where", + "compare_scalar", + "logical_not", + "index", + "sub", +) + + +class TypedInputContractTest(unittest.TestCase): + def _region(self, text: str, start: str, end: str) -> str: + """Source between two anchors, each required to occur exactly once.""" + self.assertEqual(text.count(start), 1, f"anchor not unique: {start!r}") + self.assertEqual(text.count(end), 1, f"anchor not unique: {end!r}") + begin = text.index(start) + stop = text.index(end) + self.assertLess(begin, stop, f"{start!r} must precede {end!r}") + return text[begin:stop] + + def _input_loop(self) -> str: + return self._region(_DRIVER.read_text(), _INPUT_LOOP_START, _INPUT_LOOP_END) + + def _function_source(self, path: Path, name: str) -> str: + source = path.read_text() + for node in ast.walk(ast.parse(source)): + if isinstance(node, ast.FunctionDef) and node.name == name: + segment = ast.get_source_segment(source, node) + self.assertIsNotNone(segment, f"no source for {name} in {path}") + return segment or "" + self.fail(f"{path} defines no function named {name!r}") + + def test_driver_keeps_all_three_typed_input_branches(self) -> None: + loop = self._input_loop() + for branch in (_BOOL_BRANCH, _INT32_BRANCH, _FP32_BRANCH): + self.assertEqual(loop.count(branch), 1, f"missing branch: {branch!r}") + self.assertLess(loop.index(_BOOL_BRANCH), loop.index(_INT32_BRANCH)) + self.assertLess(loop.index(_INT32_BRANCH), loop.index(_FP32_BRANCH)) + + def test_driver_loads_bool_inputs_as_int8_scalartype_bool(self) -> None: + branch = self._region(self._input_loop(), _BOOL_BRANCH, _INT32_BRANCH) + self.assertIn("load_int8_bin(in.path, n)", branch) + self.assertIn("executorch::aten::ScalarType::Bool", branch) + # A bool input widened to fp32 (or narrowed to int32) is the mutation. + self.assertNotIn("load_fp32_bin", branch) + self.assertNotIn("load_int32_bin", branch) + + def test_driver_loads_int32_inputs_as_int32(self) -> None: + branch = self._region(self._input_loop(), _INT32_BRANCH, _FP32_BRANCH) + self.assertIn("load_int32_bin(in.path, n)", branch) + self.assertNotIn("load_int8_bin", branch) + self.assertNotIn("load_fp32_bin", branch) + + def test_driver_falls_back_to_fp32_only_for_untyped_inputs(self) -> None: + loop = self._input_loop() + branch = loop[loop.index(_FP32_BRANCH) :] + self.assertIn("load_fp32_bin(in.path, n)", branch) + self.assertNotIn("ScalarType::Bool", branch) + self.assertNotIn("load_int8_bin", branch) + + def test_generator_emits_exactly_the_driver_input_dtypes(self) -> None: + write_loop = self._region( + _GENERATOR.read_text(), _WRITE_LOOP_START, _WRITE_LOOP_END + ) + emitted = re.findall(r'in_dtype = "(\w+)"', write_loop) + self.assertEqual(emitted, ["bool", "int32", "float32"]) + consumed = set(re.findall(r'in\.dtype == "(\w+)"', self._input_loop())) + # The driver names its two typed branches; fp32 is the unnamed fallback. + self.assertEqual(consumed | {"float32"}, set(emitted)) + + def test_generator_writes_bool_inputs_as_int8_not_fp32(self) -> None: + write_loop = self._region( + _GENERATOR.read_text(), _WRITE_LOOP_START, _WRITE_LOOP_END + ) + branch = self._region( + write_loop, "if t.dtype == torch.bool:", "elif t.dtype == torch.int32:" + ) + self.assertIn("_write_int8", branch) + self.assertNotIn("_write_fp32", branch) + + def test_generator_preserves_int32_golden_width(self) -> None: + source = self._function_source(_GENERATOR, "_write_golden_output") + start = source.index("elif raw.dtype == torch.int32:") + stop = source.index("\n else:", start) + branch = source[start:stop] + self.assertIn('out_dtype = "int32"', branch) + self.assertNotIn("to(torch.int64)", branch) + + def test_driver_compares_int32_outputs_exactly(self) -> None: + source = _DRIVER.read_text() + branch = self._region( + source, + '} else if (e_.golden.dtype == "int32") {', + '} else if (e_.golden.dtype == "int64") {', + ) + self.assertIn("load_int32_bin", branch) + self.assertIn("const_data_ptr()", branch) + self.assertIn("ScalarType::Int", branch) + self.assertNotIn("within_tol", branch) + + def test_generator_materializes_bool_inputs_unchanged(self) -> None: + source = self._function_source(_GENERATOR, "_materialize") + materialize_bool = self._region( + source, "if _t.dtype == torch.bool:", "return (" + ) + self.assertIn("return _t", materialize_bool) + self.assertNotIn("to(torch.int32)", materialize_bool) + self.assertNotIn("to(torch.float32)", materialize_bool) + + def test_cases_wire_one_input_generator_per_manifest_dtype(self) -> None: + cases = _CASES.read_text() + for dtype, (gen, path, literal) in _DTYPE_GENERATORS.items(): + with self.subTest(dtype=dtype): + self.assertIn(f"gen={gen}", cases, f"no {dtype} input generator wired") + self.assertIn(literal, self._function_source(path, gen)) + + def test_cases_register_every_d10_manifest_suite(self) -> None: + cases = _CASES.read_text() + for op in _D10_SUITES: + with self.subTest(op=op): + self.assertEqual(cases.count(f'@register_op_test("{op}")'), 1) diff --git a/backends/webgpu/test/ops/index/test_index.py b/backends/webgpu/test/ops/index/test_index.py index b5e83dcf329..2276dc215cf 100644 --- a/backends/webgpu/test/ops/index/test_index.py +++ b/backends/webgpu/test/ops/index/test_index.py @@ -15,6 +15,8 @@ `index_.golden.bin` so the native `test_index` self-discovers them. """ +from __future__ import annotations + import os import unittest @@ -37,11 +39,23 @@ def forward(self, x: torch.Tensor, idx: torch.Tensor) -> torch.Tensor: return x[idx] +def index_self_gen(shape: tuple[int, ...]) -> torch.Tensor: + """Distinct self values so a wrong-index gather is visible.""" + return torch.arange(shape[0], dtype=torch.float32) * 3.0 + 0.5 + + +def index_idx_gen(index_values): + """Index generator for one config's index list (int64, downcast at export).""" + + def gen(shape: tuple[int, ...]) -> torch.Tensor: + return torch.tensor(index_values, dtype=torch.int64).reshape(shape) + + return gen + + def _inputs(self_len, index_values): - # Distinct self values so a wrong-index gather is visible. - x = torch.arange(self_len, dtype=torch.float32) * 3.0 + 0.5 - idx = torch.tensor(index_values, dtype=torch.int64) - return x, idx + idx_shape = (len(index_values),) + return index_self_gen((self_len,)), index_idx_gen(index_values)(idx_shape) def _lower(x, idx): diff --git a/backends/webgpu/test/ops/scatter/__init__.py b/backends/webgpu/test/ops/scatter/__init__.py new file mode 100644 index 00000000000..2e41cd717f6 --- /dev/null +++ b/backends/webgpu/test/ops/scatter/__init__.py @@ -0,0 +1,5 @@ +# 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. diff --git a/backends/webgpu/test/ops/scatter/export_scatter_artifacts.py b/backends/webgpu/test/ops/scatter/export_scatter_artifacts.py new file mode 100644 index 00000000000..5ab83214d44 --- /dev/null +++ b/backends/webgpu/test/ops/scatter/export_scatter_artifacts.py @@ -0,0 +1,88 @@ +# 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. + +# pyre-strict + +"""Write the `webgpu_scatter_test` fixture corpus from the CPU authority.""" + +from __future__ import annotations + +import struct +import sys + +from pathlib import Path +from typing import Sequence + +from executorch.backends.webgpu.test.ops.scatter.test_scatter import ( + base_row, + destinations_are_pairwise_distinct, + EQUIVALENT_CASES, + has_official_provenance, + PROVENANCE_CASES, + scatter_cases, + scatter_parallel, + scatter_serial, + SELECTED_COUNT, +) + +BASE_FIXTURE = "base.bin" +CASES_MANIFEST = "cases.txt" + + +def _write_f32(path: Path, values: Sequence[float]) -> None: + path.write_bytes(struct.pack(f"<{len(values)}f", *values)) + + +def _write_i32(path: Path, values: Sequence[int]) -> None: + path.write_bytes(struct.pack(f"<{len(values)}i", *values)) + + +def export_scatter_artifacts(output_dir: Path) -> list[str]: + output_dir.mkdir(parents=True, exist_ok=True) + base = base_row() + _write_f32(output_dir / BASE_FIXTURE, base) + + cases = scatter_cases() + lines: list[str] = [] + for name in sorted(cases): + case = cases[name] + indices = case["indices"] + source = case["source"] + assert isinstance(indices, list) and isinstance(source, list) + if len(indices) != SELECTED_COUNT or len(source) != SELECTED_COUNT: + raise ValueError(f"scatter case {name} has the wrong width") + equivalent = destinations_are_pairwise_distinct(indices) + provenance = has_official_provenance(indices) + if equivalent != (name in EQUIVALENT_CASES): + raise ValueError( + f"scatter case {name} mislabels parallel-route equivalence" + ) + if provenance != (name in PROVENANCE_CASES): + raise ValueError(f"scatter case {name} mislabels official provenance") + expected = scatter_serial(base, indices, source) + if (scatter_parallel(base, indices, source) == expected) != equivalent: + raise ValueError( + f"scatter case {name} disagrees with its parallel-route label" + ) + _write_i32(output_dir / f"{name}.index.bin", indices) + _write_f32(output_dir / f"{name}.source.bin", source) + _write_f32(output_dir / f"{name}.expected.bin", expected) + lines.append(f"{name} {1 if equivalent else 0} {1 if provenance else 0}") + + (output_dir / CASES_MANIFEST).write_text("\n".join(lines) + "\n", encoding="utf-8") + return [line.split(" ", 1)[0] for line in lines] + + +def main(argv: Sequence[str] | None = None) -> int: + args = list(sys.argv[1:] if argv is None else argv) + if len(args) != 1: + raise SystemExit("usage: export_scatter_artifacts.py ") + export_scatter_artifacts(Path(args[0])) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/backends/webgpu/test/ops/scatter/test_scatter.py b/backends/webgpu/test/ops/scatter/test_scatter.py new file mode 100644 index 00000000000..58f86c45204 --- /dev/null +++ b/backends/webgpu/test/ops/scatter/test_scatter.py @@ -0,0 +1,345 @@ +# 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. + +# pyre-strict + +"""CPU authority for the Gemma 4 MTP scatter routes. + +`scatter_serial` transcribes `runtime/ops/scatter/scatter.wgsl` (ascending +last-write-wins, mirroring portable `op_scatter.cpp`). +`scatter_unique_indices.wgsl` is arbitrary-write-wins and is therefore only +equivalent when the destinations are pairwise distinct — the provenance the +official top-32 token ordering guarantees and this module enforces. +""" + +from __future__ import annotations + +import unittest + +VOCAB_SIZE = 262144 +SELECTED_COUNT = 4096 +CENTROIDS = 2048 +TOKENS_PER_CENTROID = 128 +SELECTED_CENTROIDS = 32 + +_U32 = 0xFFFFFFFF + + +def scatter_serial( + base: list[float], indices: list[int], source: list[float] +) -> list[float]: + """Ascending last-write-wins with the shader's verbatim range guard.""" + if len(base) != VOCAB_SIZE: + raise ValueError(f"scatter authority requires a {VOCAB_SIZE}-wide base") + if len(indices) != SELECTED_COUNT or len(source) != SELECTED_COUNT: + raise ValueError(f"scatter authority requires {SELECTED_COUNT} writes") + out = list(base) + for i in range(SELECTED_COUNT): + destination = indices[i] + if 0 <= destination < VOCAB_SIZE: + out[destination] = source[i] + return out + + +def destinations_are_pairwise_distinct(indices: list[int]) -> bool: + """Whether the parallel route is well defined: distinct *surviving* writes.""" + live = [d for d in indices if 0 <= d < VOCAB_SIZE] + return len(set(live)) == len(live) + + +def has_official_provenance(indices: list[int]) -> bool: + """The exporter-certified contract: every destination in range and distinct.""" + return all(0 <= d < VOCAB_SIZE for d in indices) and len(set(indices)) == len( + indices + ) + + +def scatter_parallel( + base: list[float], indices: list[int], source: list[float] +) -> list[float]: + """Order-independent transcription of `scatter_unique_indices.wgsl`.""" + if len(base) != VOCAB_SIZE: + raise ValueError(f"scatter authority requires a {VOCAB_SIZE}-wide base") + if len(indices) != SELECTED_COUNT or len(source) != SELECTED_COUNT: + raise ValueError(f"scatter authority requires {SELECTED_COUNT} writes") + written: dict[int, float] = {} + # Descending: a duplicate destination keeps the lowest index, not serial's last. + for i in range(SELECTED_COUNT - 1, -1, -1): + destination = indices[i] + if 0 <= destination < VOCAB_SIZE: + written[destination] = source[i] + out = list(base) + for destination, value in written.items(): + out[destination] = value + return out + + +def scatter_unique( + base: list[float], indices: list[int], source: list[float] +) -> list[float]: + """The parallel route; refuses the inputs on which it is not well defined.""" + if not destinations_are_pairwise_distinct(indices): + raise ValueError("scatter_src_unique requires pairwise-distinct destinations") + return scatter_parallel(base, indices, source) + + +def _lcg(seed: int, count: int) -> list[int]: + state = seed & _U32 + drawn: list[int] = [] + for _ in range(count): + state = (1664525 * state + 1013904223) & _U32 + drawn.append(state) + return drawn + + +def token_ordering(seed: int = 17) -> list[list[int]]: + """A deterministic [2048, 128] permutation of range(262144).""" + tokens = list(range(VOCAB_SIZE)) + draws = _lcg(seed, VOCAB_SIZE) + for i in range(VOCAB_SIZE - 1, 0, -1): + j = draws[i] % (i + 1) + tokens[i], tokens[j] = tokens[j], tokens[i] + return [ + tokens[row * TOKENS_PER_CENTROID : (row + 1) * TOKENS_PER_CENTROID] + for row in range(CENTROIDS) + ] + + +def official_selected_indices(selected_centroids: list[int]) -> list[int]: + """`token_ordering[topk_indices]` flattened, i.e. the shipped provenance.""" + if len(selected_centroids) != SELECTED_CENTROIDS: + raise ValueError(f"expected {SELECTED_CENTROIDS} selected centroids") + if len(set(selected_centroids)) != SELECTED_CENTROIDS: + raise ValueError("selected centroids must be distinct") + ordering = token_ordering() + flattened: list[int] = [] + for row in selected_centroids: + if not 0 <= row < CENTROIDS: + raise ValueError(f"centroid row {row} out of range") + flattened.extend(ordering[row]) + return flattened + + +def base_row() -> list[float]: + return [float((index % 1021) - 510) * 0.5 for index in range(VOCAB_SIZE)] + + +def source_row(seed: int) -> list[float]: + return [(draw / 2**32) * 20.0 - 10.0 for draw in _lcg(seed, SELECTED_COUNT)] + + +def _selected_centroids() -> list[int]: + return [row * 61 % CENTROIDS for row in range(SELECTED_CENTROIDS)] + + +def _differing_positions(left: list[float], right: list[float]) -> list[int]: + return [index for index in range(VOCAB_SIZE) if left[index] != right[index]] + + +def _duplicate_indices() -> list[int]: + indices = official_selected_indices(_selected_centroids()) + indices[10] = indices[4000] + indices[11] = indices[4000] + return indices + + +def _negative_indices() -> list[int]: + indices = official_selected_indices(_selected_centroids()) + indices[0] = -1 + indices[1] = -VOCAB_SIZE + return indices + + +def _out_of_range_indices() -> list[int]: + indices = official_selected_indices(_selected_centroids()) + indices[2] = VOCAB_SIZE + indices[3] = VOCAB_SIZE + 7 + return indices + + +def _boundary_indices() -> list[int]: + indices = official_selected_indices(_selected_centroids()) + indices[0] = 0 + indices[1] = VOCAB_SIZE - 1 + seen: set[int] = set() + for position, destination in enumerate(indices): + while destination in seen: + destination = (destination + 1) % VOCAB_SIZE + indices[position] = destination + seen.add(destination) + return indices + + +# Fixture contract consumed by `webgpu_scatter_test`. `equivalent` marks the cases +# on which the parallel route must match serial bit-for-bit; `provenance` marks the +# strictly smaller set an exporter may certify for `et_vk.scatter_src_unique`. +def scatter_cases() -> dict[str, dict[str, object]]: + official = official_selected_indices(_selected_centroids()) + return { + "official_unique": { + "indices": official, + "source": source_row(5), + }, + "boundary_unique": { + "indices": _boundary_indices(), + "source": source_row(19), + }, + "reversed_unique": { + "indices": list(reversed(official)), + "source": source_row(23), + }, + "duplicate_destinations": { + "indices": _duplicate_indices(), + "source": source_row(29), + }, + "negative_destinations": { + "indices": _negative_indices(), + "source": source_row(31), + }, + "out_of_range_destinations": { + "indices": _out_of_range_indices(), + "source": source_row(37), + }, + } + + +EQUIVALENT_CASES = ( + "official_unique", + "boundary_unique", + "reversed_unique", + "negative_destinations", + "out_of_range_destinations", +) +PROVENANCE_CASES = ("official_unique", "boundary_unique", "reversed_unique") + + +class TestScatterCpu(unittest.TestCase): + def test_official_provenance_is_a_full_permutation(self) -> None: + ordering = token_ordering() + self.assertEqual(len(ordering), CENTROIDS) + flattened = [token for row in ordering for token in row] + self.assertEqual(len(flattened), VOCAB_SIZE) + self.assertEqual(len(set(flattened)), VOCAB_SIZE) + + def test_official_selection_is_pairwise_distinct(self) -> None: + indices = official_selected_indices(_selected_centroids()) + self.assertEqual(len(indices), SELECTED_COUNT) + self.assertTrue(destinations_are_pairwise_distinct(indices)) + + def test_case_labels_match_the_computed_predicates(self) -> None: + cases = scatter_cases() + for name, case in cases.items(): + indices = case["indices"] + assert isinstance(indices, list) + self.assertEqual( + destinations_are_pairwise_distinct(indices), + name in EQUIVALENT_CASES, + f"{name} mislabels parallel-route equivalence", + ) + self.assertEqual( + has_official_provenance(indices), + name in PROVENANCE_CASES, + f"{name} mislabels official provenance", + ) + # Provenance is strictly stronger than equivalence: the out-of-range and + # negative rows are bit-equivalent yet must never be certified. + self.assertLess(set(PROVENANCE_CASES), set(EQUIVALENT_CASES)) + + def test_unique_route_matches_serial_on_official_indices(self) -> None: + base = base_row() + cases = scatter_cases() + for name in EQUIVALENT_CASES: + case = cases[name] + indices = case["indices"] + source = case["source"] + assert isinstance(indices, list) and isinstance(source, list) + self.assertEqual( + _differing_positions( + scatter_unique(base, indices, source), + scatter_serial(base, indices, source), + ), + [], + name, + ) + + def test_parallel_route_disagrees_with_serial_on_duplicates(self) -> None: + # Why the distinctness guard exists: the two routes are only equal there. + base = base_row() + case = scatter_cases()["duplicate_destinations"] + indices = case["indices"] + source = case["source"] + assert isinstance(indices, list) and isinstance(source, list) + parallel = scatter_parallel(base, indices, source) + serial = scatter_serial(base, indices, source) + duplicated = indices[4000] + self.assertEqual([indices[10], indices[11]], [duplicated, duplicated]) + self.assertNotEqual(source[10], source[4000]) + self.assertEqual(_differing_positions(parallel, serial), [duplicated]) + self.assertEqual(parallel[duplicated], source[10]) + self.assertEqual(serial[duplicated], source[4000]) + + def test_unique_route_rejects_duplicate_destinations(self) -> None: + base = base_row() + case = scatter_cases()["duplicate_destinations"] + indices = case["indices"] + source = case["source"] + assert isinstance(indices, list) and isinstance(source, list) + self.assertFalse(destinations_are_pairwise_distinct(indices)) + with self.assertRaisesRegex(ValueError, "pairwise-distinct"): + scatter_unique(base, indices, source) + # The duplicate-safe route stays defined: the later write wins. + result = scatter_serial(base, indices, source) + self.assertEqual(result[indices[4000]], source[4000]) + + def test_out_of_range_and_negative_destinations_are_dropped(self) -> None: + base = base_row() + for name in ("negative_destinations", "out_of_range_destinations"): + case = scatter_cases()[name] + indices = case["indices"] + source = case["source"] + assert isinstance(indices, list) and isinstance(source, list) + result = scatter_serial(base, indices, source) + dropped = [ + position + for position, destination in enumerate(indices) + if not 0 <= destination < VOCAB_SIZE + ] + self.assertTrue(dropped, name) + written = { + indices[position] + for position in range(SELECTED_COUNT) + if position not in set(dropped) + } + for index in range(VOCAB_SIZE): + if index not in written: + self.assertEqual(result[index], base[index], f"{name}@{index}") + + def test_boundary_destinations_are_written(self) -> None: + base = base_row() + case = scatter_cases()["boundary_unique"] + indices = case["indices"] + source = case["source"] + assert isinstance(indices, list) and isinstance(source, list) + result = scatter_unique(base, indices, source) + self.assertEqual(result[0], source[indices.index(0)]) + self.assertEqual(result[VOCAB_SIZE - 1], source[indices.index(VOCAB_SIZE - 1)]) + + def test_rejects_wrong_widths(self) -> None: + base = base_row() + indices = official_selected_indices(_selected_centroids()) + source = source_row(5) + with self.assertRaisesRegex(ValueError, str(VOCAB_SIZE)): + scatter_serial(base[:-1], indices, source) + with self.assertRaisesRegex(ValueError, str(SELECTED_COUNT)): + scatter_serial(base, indices[:-1], source) + with self.assertRaisesRegex(ValueError, str(SELECTED_COUNT)): + scatter_serial(base, indices, source[:-1]) + + def test_rejects_non_distinct_centroid_selection(self) -> None: + rows = _selected_centroids() + rows[1] = rows[0] + with self.assertRaisesRegex(ValueError, "distinct"): + official_selected_indices(rows) diff --git a/backends/webgpu/test/ops/test_gather.py b/backends/webgpu/test/ops/test_gather.py index 3aff9b8626f..20d1d2671f3 100644 --- a/backends/webgpu/test/ops/test_gather.py +++ b/backends/webgpu/test/ops/test_gather.py @@ -15,6 +15,7 @@ from __future__ import annotations +import math import os import unittest @@ -40,18 +41,24 @@ def forward(self, x: torch.Tensor, index: torch.Tensor) -> torch.Tensor: return torch.gather(x, self.dim, index) +def gather_self_gen(shape: tuple[int, ...]) -> torch.Tensor: + """Distinct fp32 source values, so a wrong pick is visible.""" + return torch.arange(math.prod(shape), dtype=torch.float32).reshape(shape) + + +def gather_index_gen(bound: int): + """Index generator for a given self.size(dim); cycles so rows differ.""" + + def gen(shape: tuple[int, ...]) -> torch.Tensor: + m = math.prod(shape) + return (torch.arange(m, dtype=torch.int64) % bound).reshape(shape) + + return gen + + def _det_inputs(self_shape, dim: int, index_shape): """Distinct fp32 source (a wrong pick is visible) + in-range int64 index.""" - n = 1 - for s in self_shape: - n *= s - x = torch.arange(n, dtype=torch.float32).reshape(self_shape) - bound = self_shape[dim] - m = 1 - for s in index_shape: - m *= s - index = (torch.arange(m, dtype=torch.int64) % bound).reshape(index_shape) - return x, index + return gather_self_gen(self_shape), gather_index_gen(self_shape[dim])(index_shape) def _lower(m: torch.nn.Module, x: torch.Tensor, index: torch.Tensor): diff --git a/backends/webgpu/test/ops/test_sub.py b/backends/webgpu/test/ops/test_sub.py index 5966998bd7b..bca8b0ed40c 100644 --- a/backends/webgpu/test/ops/test_sub.py +++ b/backends/webgpu/test/ops/test_sub.py @@ -32,6 +32,40 @@ } +def int32_wrap_input_a(shape): + values = torch.tensor( + [ + -(1 << 31), + (1 << 31) - 1, + 0, + -1, + 17, + -29, + 1 << 30, + -(1 << 30), + ], + dtype=torch.int32, + ) + return values[: torch.Size(shape).numel()].reshape(shape) + + +def int32_wrap_input_b(shape): + values = torch.tensor( + [ + 1, + -1, + (1 << 31) - 1, + -(1 << 31), + -31, + 37, + -(1 << 30), + 1 << 30, + ], + dtype=torch.int32, + ) + return values[: torch.Size(shape).numel()].reshape(shape) + + class SubModule(torch.nn.Module): def forward(self, a: torch.Tensor, b: torch.Tensor) -> torch.Tensor: return torch.sub(a, b) @@ -95,6 +129,54 @@ def test_golden_matches_fp64_alpha(self) -> None: ref = (a.double() - alpha * b.double()).to(torch.float32) torch.testing.assert_close(SubAlphaModule(alpha)(a, b), ref) + def test_int32_wraps_at_signed_boundaries(self) -> None: + a = int32_wrap_input_a((2, 4)) + b = int32_wrap_input_b((2, 4)) + expected = torch.tensor( + [ + (1 << 31) - 1, + -(1 << 31), + -((1 << 31) - 1), + (1 << 31) - 1, + 48, + -66, + -(1 << 31), + -(1 << 31), + ], + dtype=torch.int32, + ).reshape(2, 4) + torch.testing.assert_close(SubAlphaModule(1)(a, b), expected, rtol=0, atol=0) + + def test_int32_broadcast_is_exact(self) -> None: + a = torch.tensor( + [[-(1 << 31), 7, (1 << 31) - 1], [4, -5, 6]], + dtype=torch.int32, + ) + b = torch.tensor([[1, -3, (1 << 31) - 1]], dtype=torch.int32) + expected = torch.tensor( + [[(1 << 31) - 1, 10, 0], [3, -2, -((1 << 31) - 7)]], + dtype=torch.int32, + ) + torch.testing.assert_close(SubAlphaModule(1)(a, b), expected, rtol=0, atol=0) + + def test_int32_alpha_wraps_multiply_and_subtract(self) -> None: + a = int32_wrap_input_a((2, 4)) + b = int32_wrap_input_b((2, 4)) + expected = torch.tensor( + [ + (1 << 31) - 3, + -((1 << 31) - 2), + -((1 << 31) - 3), + (1 << 31) - 1, + 110, + -140, + 0, + 0, + ], + dtype=torch.int32, + ).reshape(2, 4) + torch.testing.assert_close(SubAlphaModule(3)(a, b), expected, rtol=0, atol=0) + def export_sub_model(pte_path: str, golden_path: str, input_path: str) -> None: """Write sub(a, b) .pte + fp64-computed torch golden + raw LE fp32 inputs (in1, in2).""" diff --git a/backends/webgpu/test/ops/test_to_copy.py b/backends/webgpu/test/ops/test_to_copy.py index 54b400ea9ef..d521144b019 100644 --- a/backends/webgpu/test/ops/test_to_copy.py +++ b/backends/webgpu/test/ops/test_to_copy.py @@ -34,6 +34,12 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return x.to(torch.int32) +class ToCopyFloatToInt64Module(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + # int64 EValue over the int32 delegate buffer (downcast_64_bit + widen). + return x.to(torch.int64) + + class ToCopyFloatToIntToFloatModule(torch.nn.Module): def forward(self, x: torch.Tensor) -> torch.Tensor: return x.to(torch.int32).to(torch.float32) diff --git a/backends/webgpu/test/ops/test_where.py b/backends/webgpu/test/ops/test_where.py index 02ccaf2d19a..bf02dcd30c8 100644 --- a/backends/webgpu/test/ops/test_where.py +++ b/backends/webgpu/test/ops/test_where.py @@ -9,11 +9,15 @@ `where(cond, a, b) -> cond ? a : b`, with cond a 1-byte bool and a/b fp32 (broadcast across all three operands). The kernel reads cond byte-packed as `array` and relinearizes each out coord onto every operand. Configs cover -the equal-shape path plus broadcasts that exercise the size-1 clamp on cond, a, -and b. The native binary has no ATen, so the golden is computed with torch here +the equal-shape path, broadcasts that exercise the size-1 clamp on cond, a, and +b, and a numel-15 tail whose cond byte buffer is not a whole number of u32 +words. The native binary has no ATen, so the golden is computed with torch here and checked in etvk CI. """ +from __future__ import annotations + +import math import unittest import torch @@ -26,6 +30,7 @@ "equal": ((4, 8), (4, 8), (4, 8)), "broadcast": ((4, 1), (4, 8), (1, 8)), "cond_row": ((8,), (4, 8), (4, 8)), + "tail": ((3, 5), (3, 5), (3, 5)), } @@ -36,13 +41,28 @@ def forward( return torch.where(cond, a, b) +def where_cond_gen(shape: tuple[int, ...]) -> torch.Tensor: + """Repeating 7-long bool mask: every u32 word sees both True and False.""" + n = math.prod(shape) + pattern = torch.tensor([1, 0, 0, 1, 1, 1, 0], dtype=torch.bool) + repeats = (n + pattern.numel() - 1) // pattern.numel() + return pattern.repeat(repeats)[:n].reshape(shape) + + +def where_a_gen(shape: tuple[int, ...]) -> torch.Tensor: + """Strictly positive; `b` is strictly negative, so a wrong pick flips sign.""" + n = math.prod(shape) + return (torch.arange(n, dtype=torch.float32) + 1.0).reshape(shape) + + +def where_b_gen(shape: tuple[int, ...]) -> torch.Tensor: + n = math.prod(shape) + return (-torch.arange(n, dtype=torch.float32) - 1.0).reshape(shape) + + def _det_inputs(cond_shape, a_shape, b_shape): """Deterministic (bool cond, fp32 a, fp32 b) for a config.""" - g = torch.Generator().manual_seed(0) - cond = torch.rand(cond_shape, generator=g) > 0.5 - a = torch.randn(*a_shape, generator=g, dtype=torch.float32) - b = torch.randn(*b_shape, generator=g, dtype=torch.float32) - return cond, a, b + return where_cond_gen(cond_shape), where_a_gen(a_shape), where_b_gen(b_shape) def _fp64_golden(cond, a, b): diff --git a/backends/webgpu/test/ops/topk/__init__.py b/backends/webgpu/test/ops/topk/__init__.py new file mode 100644 index 00000000000..2e41cd717f6 --- /dev/null +++ b/backends/webgpu/test/ops/topk/__init__.py @@ -0,0 +1,5 @@ +# 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. diff --git a/backends/webgpu/test/ops/topk/export_topk_artifacts.py b/backends/webgpu/test/ops/topk/export_topk_artifacts.py new file mode 100644 index 00000000000..dcf7fe08420 --- /dev/null +++ b/backends/webgpu/test/ops/topk/export_topk_artifacts.py @@ -0,0 +1,89 @@ +# 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. + +# pyre-strict + +"""Write the `webgpu_topk_test` fixture corpus from the sealed CPU authority.""" + +from __future__ import annotations + +import json +import struct +import sys + +from pathlib import Path +from typing import Mapping, Sequence + +from executorch.backends.webgpu.test.ops.topk.test_topk import ( + authority_digest, + AUTHORITY_SHA256, + INPUT_WIDTH, + OUTPUT_WIDTH, +) + +CASES_MANIFEST = "cases.txt" + + +def _write_u32(path: Path, words: Sequence[int]) -> None: + path.write_bytes(struct.pack(f"<{len(words)}I", *words)) + + +def _write_i32(path: Path, values: Sequence[int]) -> None: + path.write_bytes(struct.pack(f"<{len(values)}i", *values)) + + +def _load_sealed_authority(authority_path: Path) -> Mapping[str, object]: + sealed = json.loads(authority_path.read_text(encoding="utf-8")) + digest = authority_digest(sealed) + if digest != AUTHORITY_SHA256 or sealed.get("sha256") != AUTHORITY_SHA256: + raise ValueError( + f"top-k authority {authority_path} does not match the committed " + f"digest {AUTHORITY_SHA256} (body {digest}, seal " + f"{sealed.get('sha256')!r}); a deliberate authority change has to " + "update AUTHORITY_SHA256 in test_topk.py" + ) + return sealed + + +def export_topk_artifacts(output_dir: Path, authority_path: Path) -> list[str]: + authority = _load_sealed_authority(authority_path) + cases = authority["cases"] + assert isinstance(cases, dict) + + output_dir.mkdir(parents=True, exist_ok=True) + names = sorted(cases) + for name in names: + case = cases[name] + assert isinstance(case, dict) + scores = case["scores_bits"] + values = case["values_bits"] + indices = case["indices"] + assert isinstance(scores, list) and isinstance(values, list) + assert isinstance(indices, list) + if len(scores) != INPUT_WIDTH or len(values) != OUTPUT_WIDTH: + raise ValueError(f"top-k authority case {name} has the wrong width") + if len(indices) != OUTPUT_WIDTH: + raise ValueError(f"top-k authority case {name} has the wrong index width") + _write_u32(output_dir / f"{name}.scores.bin", scores) + _write_u32(output_dir / f"{name}.values.bin", values) + _write_i32(output_dir / f"{name}.indices.bin", indices) + + (output_dir / CASES_MANIFEST).write_text("\n".join(names) + "\n", encoding="utf-8") + return names + + +def main(argv: Sequence[str] | None = None) -> int: + args = list(sys.argv[1:] if argv is None else argv) + if len(args) != 2: + raise SystemExit( + "usage: export_topk_artifacts.py " + ) + export_topk_artifacts(Path(args[0]), Path(args[1])) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/backends/webgpu/test/ops/topk/test_topk.py b/backends/webgpu/test/ops/topk/test_topk.py new file mode 100644 index 00000000000..afbdfab0726 --- /dev/null +++ b/backends/webgpu/test/ops/topk/test_topk.py @@ -0,0 +1,412 @@ +# 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. + +# pyre-strict + +"""CPU authority for the Gemma 4 MTP top-k route. + +`topk_reference` transcribes `runtime/ops/topk/topk.wgsl` bit-for-bit: the same +32-entry heap, the same bit-pattern comparator, and the same emission order. It +is the oracle the WG64-staged shader must reproduce exactly, so it deliberately +does NOT adopt `torch.topk` tie behaviour. + +`AUTHORITY_SHA256` pins the exported corpus: any change to the transcription or +to `topk_cases()` moves the digest and has to be re-committed on purpose. +""" + +from __future__ import annotations + +import hashlib +import json +import os +import struct +import unittest + +from pathlib import Path +from typing import Mapping + +INPUT_WIDTH = 2048 +OUTPUT_WIDTH = 32 + +# Committed digest of the canonical authority payload; see the module docstring. +AUTHORITY_SHA256 = "c7289f314f7c364251ae228f0c28cc300677db333c649c5a99c39994a259d50d" + +_U32 = 0xFFFFFFFF + + +def f32_bits(value: float) -> int: + return struct.unpack(" float: + return struct.unpack(" bool: + return (bits & 0x7F800000) == 0x7F800000 and (bits & 0x007FFFFF) != 0 + + +def float_less_than_bits(lhs: int, rhs: int) -> bool: + lhs_nan = _is_nan_bits(lhs) + rhs_nan = _is_nan_bits(rhs) + if lhs_nan or rhs_nan: + return (not lhs_nan) and rhs_nan + + lhs_magnitude = lhs & 0x7FFFFFFF + rhs_magnitude = rhs & 0x7FFFFFFF + if lhs_magnitude == 0 and rhs_magnitude == 0: + return False + + lhs_negative = (lhs & 0x80000000) != 0 + rhs_negative = (rhs & 0x80000000) != 0 + if lhs_negative != rhs_negative: + return lhs_negative + if lhs_negative: + return lhs > rhs + return lhs < rhs + + +def _greater(lhs: int, rhs: int) -> bool: + return float_less_than_bits(rhs, lhs) + + +def _push_heap( + heap_values: list[int], + heap_indices: list[int], + initial_hole: int, + top: int, + value_bits: int, + value_index: int, +) -> None: + hole = initial_hole + while hole > top: + parent = (hole - 1) // 2 + if not _greater(heap_values[parent], value_bits): + break + heap_values[hole] = heap_values[parent] + heap_indices[hole] = heap_indices[parent] + hole = parent + heap_values[hole] = value_bits + heap_indices[hole] = value_index + + +def _adjust_heap( + heap_values: list[int], + heap_indices: list[int], + initial_hole: int, + length: int, + value_bits: int, + value_index: int, +) -> None: + top = initial_hole + hole = initial_hole + second_child = initial_hole + while second_child < (length - 1) // 2: + second_child = 2 * (second_child + 1) + if _greater(heap_values[second_child], heap_values[second_child - 1]): + second_child -= 1 + heap_values[hole] = heap_values[second_child] + heap_indices[hole] = heap_indices[second_child] + hole = second_child + if (length & 1) == 0 and second_child == (length - 2) // 2: + second_child = 2 * (second_child + 1) + heap_values[hole] = heap_values[second_child - 1] + heap_indices[hole] = heap_indices[second_child - 1] + hole = second_child - 1 + _push_heap(heap_values, heap_indices, hole, top, value_bits, value_index) + + +def topk_reference(scores_bits: list[int]) -> tuple[list[int], list[int]]: + """Return (values_bits, indices) exactly as `topk.wgsl` emits them.""" + if len(scores_bits) != INPUT_WIDTH: + raise ValueError(f"topk authority requires {INPUT_WIDTH} scores") + + heap_values = list(scores_bits[:OUTPUT_WIDTH]) + heap_indices = list(range(OUTPUT_WIDTH)) + + for parent in range(OUTPUT_WIDTH // 2 - 1, -1, -1): + _adjust_heap( + heap_values, + heap_indices, + parent, + OUTPUT_WIDTH, + heap_values[parent], + heap_indices[parent], + ) + + heap_root = heap_values[0] + for index in range(OUTPUT_WIDTH, INPUT_WIDTH): + value_bits = scores_bits[index] + if _greater(value_bits, heap_root): + _adjust_heap(heap_values, heap_indices, 0, OUTPUT_WIDTH, value_bits, index) + heap_root = heap_values[0] + + last = OUTPUT_WIDTH + while last > 1: + last -= 1 + value_bits = heap_values[last] + value_index = heap_indices[last] + heap_values[last] = heap_values[0] + heap_indices[last] = heap_indices[0] + _adjust_heap(heap_values, heap_indices, 0, last, value_bits, value_index) + + return heap_values, heap_indices + + +def _lcg(seed: int, count: int) -> list[int]: + state = seed & _U32 + drawn: list[int] = [] + for _ in range(count): + state = (1664525 * state + 1013904223) & _U32 + drawn.append(state) + return drawn + + +def _ramp_row() -> list[int]: + return [f32_bits(-1024.0 + index * 1.0) for index in range(INPUT_WIDTH)] + + +def _random_row(seed: int) -> list[int]: + return [ + f32_bits((draw / 2**32) * 200.0 - 100.0) for draw in _lcg(seed, INPUT_WIDTH) + ] + + +def _interior_ties_row() -> list[int]: + row = _random_row(7) + for index in range(200, 264): + row[index] = f32_bits(50.0) + return row + + +def _boundary_ties_row() -> list[int]: + row = [f32_bits(-1.0)] * INPUT_WIDTH + for index in range(0, 40): + row[index * 7] = f32_bits(3.5) + return row + + +def _all_equal_row() -> list[int]: + return [f32_bits(2.25)] * INPUT_WIDTH + + +def _nan_row() -> list[int]: + row = _random_row(11) + row[5] = 0x7FC00001 + row[600] = 0x7F800001 + row[1900] = 0xFFC00007 + return row + + +def _signed_zero_row() -> list[int]: + row = [f32_bits(-3.0)] * INPUT_WIDTH + for index in range(0, 64, 2): + row[index] = 0x00000000 + row[index + 1] = 0x80000000 + return row + + +def _infinity_row() -> list[int]: + row = _random_row(13) + row[0] = 0xFF800000 + row[1] = 0x7F800000 + row[1000] = 0x7F800000 + row[2047] = 0xFF800000 + return row + + +def _descending_row() -> list[int]: + return [f32_bits(float(INPUT_WIDTH - index)) for index in range(INPUT_WIDTH)] + + +# The only corpus rows whose emitted top-32 exercises the `lhs > rhs` branch. +def _all_negative_row() -> list[int]: + return [f32_bits(-1.0 - (draw / 2**32) * 100.0) for draw in _lcg(23, INPUT_WIDTH)] + + +def _straddling_zero_row() -> list[int]: + row = _all_negative_row() + for slot, index in enumerate(range(31, INPUT_WIDTH, 173)): + row[index] = f32_bits(0.25 + slot * 0.5) + return row + + +# `export_topk_artifacts.py` writes these names sorted; that sorted list is the +# `kCases[]` contract in `test/native/test_topk.cpp`. +def topk_cases() -> dict[str, list[int]]: + return { + "ordinary": _ramp_row(), + "random_seeded": _random_row(3), + "interior_ties": _interior_ties_row(), + "boundary_ties": _boundary_ties_row(), + "all_equal": _all_equal_row(), + "nan_payloads": _nan_row(), + "signed_zeros": _signed_zero_row(), + "infinities": _infinity_row(), + "descending": _descending_row(), + "all_negative": _all_negative_row(), + "straddling_zero": _straddling_zero_row(), + } + + +def authority_digest(body: Mapping[str, object]) -> str: + """SHA-256 of the canonical authority payload, excluding the seal itself.""" + payload = {key: value for key, value in body.items() if key != "sha256"} + canonical = json.dumps(payload, sort_keys=True, separators=(",", ":")) + return hashlib.sha256(canonical.encode("utf-8")).hexdigest() + + +def build_authority() -> dict[str, object]: + cases: dict[str, object] = {} + for name, scores in topk_cases().items(): + values, indices = topk_reference(scores) + cases[name] = { + "indices": indices, + "scores_bits": scores, + "values_bits": values, + } + body: dict[str, object] = { + "cases": cases, + "input_width": INPUT_WIDTH, + "output_width": OUTPUT_WIDTH, + "schema_version": 1, + } + body["sha256"] = authority_digest(body) + return body + + +class TestEagleTopKCpu(unittest.TestCase): + def test_eager_reference_is_repeatable(self) -> None: + first = build_authority() + second = build_authority() + self.assertEqual(first, second) + + for name, case in first["cases"].items(): + self.assertEqual(len(case["values_bits"]), OUTPUT_WIDTH, name) + self.assertEqual(len(case["indices"]), OUTPUT_WIDTH, name) + self.assertEqual(len(set(case["indices"])), OUTPUT_WIDTH, name) + for index in case["indices"]: + self.assertTrue(0 <= index < INPUT_WIDTH, name) + for slot, index in enumerate(case["indices"]): + self.assertEqual( + case["values_bits"][slot], case["scores_bits"][index], name + ) + + receipt = os.environ.get("EAGLE_TOPK_EAGER_RECEIPT") + if receipt: + path = Path(receipt) + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + json.dumps(first, indent=2, sort_keys=True) + "\n", encoding="utf-8" + ) + + def test_authority_matches_the_committed_digest(self) -> None: + authority = build_authority() + self.assertEqual(authority["sha256"], AUTHORITY_SHA256) + self.assertEqual(authority_digest(authority), AUTHORITY_SHA256) + + def test_emission_order_is_descending(self) -> None: + for name, scores in topk_cases().items(): + values, _ = topk_reference(scores) + for slot in range(1, OUTPUT_WIDTH): + self.assertFalse( + float_less_than_bits(values[slot - 1], values[slot]), + f"{name} slot {slot} is not in descending order", + ) + + def test_matches_sorted_selection_on_distinct_values(self) -> None: + # With all-distinct finite values the heap order is the unique answer, so + # a plain sort is an independent cross-check of the transcription. + shipped = topk_cases() + names = ( + "ordinary", + "random_seeded", + "descending", + "all_negative", + "straddling_zero", + ) + for name in names: + scores = shipped[name] + self.assertEqual(len(set(scores)), INPUT_WIDTH, name) + expected = sorted( + range(INPUT_WIDTH), key=lambda i: bits_f32(scores[i]), reverse=True + )[:OUTPUT_WIDTH] + _, indices = topk_reference(scores) + self.assertEqual(indices, expected, name) + + def test_negative_rows_keep_negatives_inside_the_selected_top_k(self) -> None: + # Without this the exported corpus never reaches `topk.wgsl:62`. + shipped = topk_cases() + negatives = shipped["all_negative"] + _, indices = topk_reference(negatives) + self.assertTrue(all(bits_f32(negatives[i]) < 0.0 for i in indices)) + + straddling = shipped["straddling_zero"] + _, indices = topk_reference(straddling) + selected = [bits_f32(straddling[i]) for i in indices] + self.assertEqual(len([value for value in selected if value > 0.0]), 12) + self.assertEqual(len([value for value in selected if value < 0.0]), 20) + + def test_comparator_ranks_nan_above_every_non_nan(self) -> None: + for nan in (0x7FC00001, 0x7F800001, 0xFFC00007): + for other in (f32_bits(3.5), f32_bits(-3.5), 0x7F800000, 0xFF800000): + self.assertTrue(float_less_than_bits(other, nan)) + self.assertFalse(float_less_than_bits(nan, other)) + + def test_comparator_treats_any_two_nans_as_equal(self) -> None: + nans = (0x7FC00001, 0x7F800001, 0xFFC00007) + for lhs in nans: + for rhs in nans: + self.assertFalse(float_less_than_bits(lhs, rhs)) + + def test_comparator_treats_signed_zeros_as_equal(self) -> None: + self.assertFalse(float_less_than_bits(0x80000000, 0x00000000)) + self.assertFalse(float_less_than_bits(0x00000000, 0x80000000)) + + def test_comparator_ranks_every_negative_below_every_positive(self) -> None: + for negative in (f32_bits(-1e-30), f32_bits(-1.0), 0xFF800000): + for positive in (f32_bits(1e-30), f32_bits(1.0), 0x7F800000): + self.assertTrue(float_less_than_bits(negative, positive)) + self.assertFalse(float_less_than_bits(positive, negative)) + + def test_comparator_orders_two_negatives_by_magnitude(self) -> None: + self.assertTrue(float_less_than_bits(f32_bits(-2.0), f32_bits(-1.0))) + self.assertFalse(float_less_than_bits(f32_bits(-1.0), f32_bits(-2.0))) + self.assertTrue(float_less_than_bits(0xFF800000, f32_bits(-3.4e38))) + + def test_comparator_orders_two_positives_by_magnitude(self) -> None: + self.assertTrue(float_less_than_bits(f32_bits(1.0), f32_bits(2.0))) + self.assertFalse(float_less_than_bits(f32_bits(2.0), f32_bits(1.0))) + self.assertTrue(float_less_than_bits(f32_bits(3.4e38), 0x7F800000)) + + def test_nan_sorts_above_every_finite_value(self) -> None: + values, indices = topk_reference(_nan_row()) + self.assertEqual(indices[:3], [5, 600, 1900]) + for bits in values[:3]: + self.assertTrue(_is_nan_bits(bits)) + self.assertFalse(_is_nan_bits(values[3])) + + def test_positive_infinity_beats_finite_and_negative_infinity_loses(self) -> None: + _, indices = topk_reference(_infinity_row()) + self.assertIn(1, indices) + self.assertIn(1000, indices) + self.assertNotIn(0, indices) + self.assertNotIn(2047, indices) + + def test_signed_zeros_compare_equal_and_beat_negatives(self) -> None: + values, indices = topk_reference(_signed_zero_row()) + self.assertEqual(sorted(indices), list(range(OUTPUT_WIDTH))) + for bits in values: + self.assertIn(bits, (0x00000000, 0x80000000)) + + def test_tie_rule_is_not_lowest_index_wins(self) -> None: + # `tie_by_low_index` is an explicitly killed mutation of this authority. + _, indices = topk_reference(_all_equal_row()) + self.assertNotEqual(indices, list(range(OUTPUT_WIDTH))) + + def test_rejects_wrong_input_width(self) -> None: + with self.assertRaisesRegex(ValueError, str(INPUT_WIDTH)): + topk_reference([f32_bits(0.0)] * (INPUT_WIDTH - 1)) diff --git a/backends/webgpu/test/targets.bzl b/backends/webgpu/test/targets.bzl index ecb4e086eeb..91d31e024d7 100644 --- a/backends/webgpu/test/targets.bzl +++ b/backends/webgpu/test/targets.bzl @@ -9,7 +9,7 @@ def define_common_targets(is_fbcode = False): python_unittest( name = "test_add", srcs = [ - "ops/add/test_add.py", + "ops/test_add.py", ], deps = [ "//caffe2:torch", @@ -44,3 +44,9 @@ def define_common_targets(is_fbcode = False): "//executorch/backends/vulkan:vulkan_preprocess", ], ) + + runtime.python_library( + name = "test_native_ci_contract", + srcs = ["test_native_ci_contract.py"], + typing = True, + ) diff --git a/backends/webgpu/test/test_native_ci_contract.py b/backends/webgpu/test/test_native_ci_contract.py index 49e81227f55..2cff6568a87 100644 --- a/backends/webgpu/test/test_native_ci_contract.py +++ b/backends/webgpu/test/test_native_ci_contract.py @@ -4,11 +4,112 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. +# pyre-strict + +from __future__ import annotations + +import ast import pathlib import re import shlex +import shutil +import subprocess +import tempfile import unittest +_BACKEND: pathlib.Path = pathlib.Path(__file__).parents[1] +_EXECUTORCH: pathlib.Path = _BACKEND.parents[1] +_GEMMA4_TESTS: pathlib.Path = _EXECUTORCH / "examples/models/gemma4/tests" + +_CI_SCRIPT: pathlib.Path = _BACKEND / "scripts/test_webgpu_native_ci.sh" +_CMAKE: pathlib.Path = _BACKEND / "CMakeLists.txt" +_DYNAMIC_SHAPE_TEST: pathlib.Path = _BACKEND / "test/native/test_dynamic_shape.cpp" +_UPDATE_CACHE_TEST: pathlib.Path = _BACKEND / "test/native/test_update_cache.cpp" +_SLICE_IMPL: pathlib.Path = _BACKEND / "runtime/ops/slice/Slice.cpp" +_SLICE_DISPATCH: pathlib.Path = _BACKEND / "runtime/ops/slice/SliceDispatch.h" + +_THIS_GATE = "backends/webgpu/test/test_native_ci_contract.py" + +# Directories D10 adds or edits test sources in; the drift guard narrows source +# control to these so an unrelated working-copy edit cannot redden it. +_D10_TEST_SURFACE: tuple[str, ...] = ( + "backends/webgpu/test/", + "examples/models/gemma4/tests/", +) + +# Every test source D10 adds or edits, executorch-relative. +_D10_TEST_SOURCES: tuple[str, ...] = ( + "backends/webgpu/test/native/test_q4gsw_m3.cpp", + "backends/webgpu/test/native/test_scatter.cpp", + "backends/webgpu/test/native/test_topk.cpp", + "backends/webgpu/test/op_tests/test_typed_input_contract.py", + "backends/webgpu/test/ops/index/test_index.py", + "backends/webgpu/test/ops/scatter/test_scatter.py", + "backends/webgpu/test/ops/test_gather.py", + "backends/webgpu/test/ops/test_to_copy.py", + "backends/webgpu/test/ops/test_where.py", + "backends/webgpu/test/ops/topk/test_topk.py", + "backends/webgpu/test/test_native_ci_contract.py", + "examples/models/gemma4/tests/test_eagle_combined_round.py", + "examples/models/gemma4/tests/test_export_assistant_webgpu_artifacts.py", + "examples/models/gemma4/tests/test_export_partitioners.py", + "examples/models/gemma4/tests/test_gemma4_spec_runner_contract.cpp", + "examples/models/gemma4/tests/test_mtp_spec_oracle.py", + "examples/models/gemma4/tests/test_oss_source_closure.py", + "examples/models/gemma4/tests/test_webgpu_artifact_manifest.py", + "examples/models/gemma4/tests/test_webgpu_spec_contract.py", +) + +# Build files whose registrations must resolve to real sources on disk. +_BUILD_FILES: tuple[pathlib.Path, ...] = ( + _BACKEND / "test/BUCK", + _BACKEND / "test/targets.bzl", + _GEMMA4_TESTS / "targets.bzl", +) + +# Buck package -> every build file allowed to define that package's targets. +_PACKAGE_BUILD_FILES: dict[str, tuple[pathlib.Path, ...]] = { + "//backends/webgpu/test": ( + _BACKEND / "test/BUCK", + _BACKEND / "test/targets.bzl", + ), + "//examples/models/gemma4/tests": (_GEMMA4_TESTS / "targets.bzl",), +} + +# GEMMA4_D10_MTP_BUCK_TEST_TARGETS from the D10 command contract, verbatim. +_MTP_BUCK_TEST_TARGETS: tuple[str, ...] = ( + "//examples/models/gemma4/tests:test_eagle_combined_round", + "//examples/models/gemma4/tests:test_export_assistant_webgpu_artifacts", + "//examples/models/gemma4/tests:test_export_partitioners", + "//examples/models/gemma4/tests:test_webgpu_artifact_manifest", + "//examples/models/gemma4/tests:test_mtp_spec_oracle", + "//examples/models/gemma4/tests:test_webgpu_spec_contract", + "//backends/webgpu/test:test_scatter_cpu", + "//backends/webgpu/test:test_topk_cpu", + "//backends/webgpu/test:test_to_copy", + "//backends/webgpu/test:test_index", +) + +# GEMMA4_D10_PLAIN_REGRESSION_TARGETS from the same contract, verbatim. +_PLAIN_REGRESSION_TARGETS: tuple[str, ...] = ( + "//examples/models/gemma4/tests:test_webgpu_artifact_manifest", + "//examples/models/gemma4/tests:test_export_partitioners", + "//examples/models/gemma4/tests:test_export_smoke", + "//examples/models/gemma4/tests:test_selected_row_cross_decoder", + "//backends/webgpu/test:test_et_vk_sdpa", + "//backends/webgpu/test:test_rope_hf_single", +) + +# Executables D10's documented `cmake --build --target ...` line names. +_CMAKE_BUILD_TARGETS: tuple[str, ...] = ( + "webgpu_native_test", + "webgpu_dynamic_shape_test", + "webgpu_update_cache_test", + "webgpu_op_test", + "webgpu_scatter_test", + "webgpu_topk_test", +) + def _bash_array(source: str, name: str) -> list[str]: match = re.search( @@ -19,16 +120,634 @@ def _bash_array(source: str, name: str) -> list[str]: return shlex.split(match.group(1)) +def _bash_function(source: str, name: str) -> str: + pattern = "^" + re.escape(name) + r"\(\) \{\n(.*?)^\}" + match = re.search( + pattern, + source, + re.MULTILINE | re.DOTALL, + ) + if match is None: + raise AssertionError(f"{name} Bash function not found") + return match.group(1) + + +def _run_required_gtests( + script: str, output: str, status: int = 0 +) -> subprocess.CompletedProcess[str]: + program = f""" +run_with_required_device() {{ +{_bash_function(script, "run_with_required_device")} +}} +run_required_gtests() {{ +{_bash_function(script, "run_required_gtests")} +}} +fake_gtest() {{ + printf '%s\\n' "$1" + return "$2" +}} +run_required_gtests fake_gtest "$1" "$2" +""" + return subprocess.run( + ["bash", "-c", program, "required-gtests", output, str(status)], + capture_output=True, + text=True, + timeout=10, + ) + + +def _run_recreate_exact_directory( + script: str, target: pathlib.Path, expected: pathlib.Path +) -> subprocess.CompletedProcess[str]: + program = f""" +recreate_exact_directory() {{ +{_bash_function(script, "recreate_exact_directory")} +}} +recreate_exact_directory "$1" "$2" +""" + return subprocess.run( + ["bash", "-c", program, "recreate-exact-directory", str(target), str(expected)], + capture_output=True, + text=True, + timeout=10, + ) + + +def _ci_script() -> str: + return _CI_SCRIPT.read_text() + + +def _required_unique_binding(source: str, name: str, value: str, before: int) -> None: + pattern = re.compile( + rf"^[ \t]*(?:(?:export|readonly|declare|typeset)" + rf"(?:[ \t]+-[A-Za-z]+)?[ \t]+)?{re.escape(name)}=[^\n]*$", + re.MULTILINE, + ) + bindings = list(pattern.finditer(source)) + expected = f"{name}={value}" + if len(bindings) != 1 or bindings[0].group() != expected: + raise AssertionError( + f"expected one canonical {name} binding, got " + f"{[binding.group() for binding in bindings]}" + ) + if bindings[0].start() >= before: + raise AssertionError(f"{name} must be bound before required Slice tests") + + +def _required_gtest_invocation(source: str) -> tuple[str, ...]: + starts = list( + re.finditer(r"^[ \t]*run_required_gtests[ \t]+", source, re.MULTILINE) + ) + if len(starts) != 1: + raise AssertionError( + f"expected one run_required_gtests invocation, got {len(starts)}" + ) + _required_unique_binding( + source, "DYNAMIC_SHAPE_DIR", '"/tmp/dynamic_shape"', starts[0].start() + ) + _required_unique_binding( + source, + "BIN_DIR", + '"${BUILD_DIR}/backends/webgpu"', + starts[0].start(), + ) + + continued = re.sub(r"\\\n[ \t]*", " ", source) + commands = [ + line.lstrip() + for line in continued.splitlines() + if line.lstrip().startswith("run_required_gtests ") + ] + if len(commands) != 1: + raise AssertionError( + f"expected one run_required_gtests invocation, got {len(commands)}" + ) + program = f""" +run_required_gtests() {{ + printf '%s\\n' "$@" +}} +BIN_DIR="$1" +DYNAMIC_SHAPE_DIR="$2" +{commands[0]} +""" + completed = subprocess.run( + [ + "bash", + "-c", + program, + "required-gtest-invocation", + "/contract bin", + "/contract dynamic shape", + ], + capture_output=True, + text=True, + timeout=10, + ) + if completed.returncode != 0: + raise AssertionError(completed.stderr) + return tuple(completed.stdout.splitlines()) + + +def _required_slice_contract_result(script: str) -> unittest.TestResult: + global _CI_SCRIPT + with tempfile.TemporaryDirectory() as temporary: + mutated_script = pathlib.Path(temporary) / "test_webgpu_native_ci.sh" + mutated_script.write_text(script) + original_script = _CI_SCRIPT + try: + _CI_SCRIPT = mutated_script + result = unittest.TestResult() + TestNativeCIContract( + "test_runs_required_slice_regressions_fail_closed" + ).run(result) + return result + finally: + _CI_SCRIPT = original_script + + +def _srcs_entries(source: str) -> list[str]: + entries: list[str] = [] + for block in re.findall(r"srcs\s*=\s*\[(.*?)\]", source, re.DOTALL): + entries.extend(re.findall(r'"([^"]+)"', block)) + return entries + + +def _target_names(source: str) -> set[str]: + return set(re.findall(r'name\s*=\s*"([^"]+)"', source)) + + +def _cmake_defines(cmake: str, name: str) -> bool: + """A bare substring match lets a longer target name mask a deleted one.""" + pattern = rf"(?:add_webgpu_native_test|add_executable)\(\s*{re.escape(name)}\b" + return re.search(pattern, cmake) is not None + + +def _undefined_labels(labels: tuple[str, ...]) -> list[str]: + undefined: list[str] = [] + for label in labels: + package, _, name = label.partition(":") + defined: set[str] = set() + for build_file in _PACKAGE_BUILD_FILES[package]: + defined |= _target_names(build_file.read_text()) + if name not in defined: + undefined.append(label) + return undefined + + +def _sl_status_paths() -> list[str] | None: + """Executorch-relative added/modified paths, or None when `sl` cannot answer.""" + try: + completed = subprocess.run( + ["sl", "status", "--reason", "D10 registration contract"], + capture_output=True, + cwd=_EXECUTORCH, + text=True, + timeout=300, + ) + except (OSError, subprocess.SubprocessError): + return None + if completed.returncode != 0: + return None + return [ + line[2:] + for line in completed.stdout.splitlines() + if line[:2] in ("M ", "A ") and not line[2:].startswith("..") + ] + + +def _sl_log_node(relative: str) -> str: + """Node that last touched a path; empty while the path is uncommitted.""" + args = ["sl", "log", relative, "-T", "{node}\n", "-l", "1"] + try: + completed = subprocess.run( + args + ["--reason", "D10 registration contract"], + capture_output=True, + cwd=_EXECUTORCH, + text=True, + timeout=300, + ) + except (OSError, subprocess.SubprocessError): + return "" + if completed.returncode != 0: + return "" + lines = completed.stdout.splitlines() + return lines[0] if lines else "" + + +def _sapling_is_usable() -> bool: + """True when `sl` and a Sapling working copy are both present.""" + if shutil.which("sl") is None: + return False + return any( + (parent / ".sl").is_dir() or (parent / ".hg").is_dir() + for parent in (_EXECUTORCH, *_EXECUTORCH.parents) + ) + + +def _base_name(node: ast.expr) -> str: + if isinstance(node, ast.Attribute): + return node.attr + if isinstance(node, ast.Name): + return node.id + return "" + + +def _defines_test_case(path: pathlib.Path) -> bool: + source = path.read_text() + if path.suffix == ".cpp": + return "TEST(" in source or "TEST_F(" in source + return any( + isinstance(node, ast.ClassDef) + and any(_base_name(base).endswith("TestCase") for base in node.bases) + for node in ast.walk(ast.parse(source)) + ) + + +def _is_test_source(relative: str) -> bool: + path = pathlib.PurePosixPath(relative) + if path.suffix not in (".py", ".cpp"): + return False + return any(part in ("test", "tests") for part in path.parts[:-1]) + + +def _reported_test_sources() -> list[str] | None: + """Test sources source control reports inside D10's surface, or None.""" + reported = _sl_status_paths() + if reported is None: + return None + return sorted( + relative + for relative in reported + if relative.startswith(_D10_TEST_SURFACE) + and _is_test_source(relative) + and (_EXECUTORCH / relative).is_file() + and _defines_test_case(_EXECUTORCH / relative) + ) + + +def _invocation_text() -> tuple[str, list[str]]: + """The CI script's commands, comments dropped and continuations joined.""" + body = "\n".join( + line for line in _ci_script().splitlines() if not line.lstrip().startswith("#") + ) + joined = re.sub(r"\\\n\s*", " ", body) + return joined, re.findall(r'-c\s+"(.*?)"', joined, re.DOTALL) + + +def _resolve_module(dotted: str) -> str | None: + parts = dotted.split(".")[1:] + while parts: + candidate = pathlib.PurePosixPath(*parts).with_suffix(".py") + if (_EXECUTORCH / candidate).is_file(): + return str(candidate) + parts.pop() + return None + + +def _ci_script_modules() -> set[str]: + """Executorch-relative sources the CI script actually invokes, not mentions.""" + joined, programs = _invocation_text() + dotted_names = re.findall( + r"-m\s+(?:unittest\s+)?(executorch(?:\.[A-Za-z_]\w*)+)", joined + ) + for program in programs: + dotted_names += re.findall(r"executorch(?:\.[A-Za-z_]\w*)+", program) + invoked: set[str] = set() + for dotted in dotted_names: + resolved = _resolve_module(dotted) + if resolved is not None: + invoked.add(resolved) + return invoked + + +def _registered_sources() -> set[str]: + registered: set[str] = set() + for build_file in _BUILD_FILES: + for entry in _srcs_entries(build_file.read_text()): + resolved = (build_file.parent / entry).resolve() + registered.add(str(resolved.relative_to(_EXECUTORCH.resolve()))) + for entry in re.findall(r"[\w][\w./]*\.cpp", _CMAKE.read_text()): + registered.add(str(pathlib.PurePosixPath("backends/webgpu") / entry)) + return registered | _ci_script_modules() + + class TestNativeCIContract(unittest.TestCase): + def test_required_slice_contract_rejects_missing_heavy_env(self) -> None: + script = _ci_script() + invocation = ( + "run_required_gtests env WEBGPU_REQUIRE_DEVICE=1 " + "WEBGPU_TEST_HEAVY=1 \\\n" + ) + self.assertEqual(1, script.count(invocation)) + result = _required_slice_contract_result( + script.replace( + invocation, + "run_required_gtests env WEBGPU_REQUIRE_DEVICE=1 \\\n", + ) + ) + + self.assertFalse(result.wasSuccessful()) + self.assertEqual([], result.errors) + self.assertEqual(1, len(result.failures)) + + def test_required_slice_contract_rejects_stale_artifact_directory(self) -> None: + script = _ci_script() + invocation = ( + '"${BIN_DIR}/webgpu_dynamic_shape_test" "${DYNAMIC_SHAPE_DIR}" \\\n' + ) + self.assertEqual(1, script.count(invocation)) + result = _required_slice_contract_result( + script.replace( + invocation, + '"${BIN_DIR}/webgpu_dynamic_shape_test" ' + '"/tmp/stale_dynamic_shape" \\\n', + ) + ) + + self.assertFalse(result.wasSuccessful()) + self.assertEqual([], result.errors) + self.assertEqual(1, len(result.failures)) + + def test_required_slice_contract_rejects_late_artifact_rebinding(self) -> None: + script = _ci_script() + invocation = ( + "run_required_gtests env WEBGPU_REQUIRE_DEVICE=1 " + "WEBGPU_TEST_HEAVY=1 \\\n" + ) + self.assertEqual(1, script.count(invocation)) + result = _required_slice_contract_result( + script.replace( + invocation, + "DYNAMIC_SHAPE_DIR=/tmp/stale_dynamic_shape\n" + invocation, + ) + ) + + self.assertFalse(result.wasSuccessful()) + self.assertEqual([], result.errors) + self.assertEqual(1, len(result.failures)) + + def test_required_slice_contract_rejects_literal_variable_arguments(self) -> None: + script = _ci_script() + invocation = ( + '"${BIN_DIR}/webgpu_dynamic_shape_test" "${DYNAMIC_SHAPE_DIR}" \\\n' + ) + self.assertEqual(1, script.count(invocation)) + result = _required_slice_contract_result( + script.replace( + invocation, + "'${BIN_DIR}/webgpu_dynamic_shape_test' " "'${DYNAMIC_SHAPE_DIR}' \\\n", + ) + ) + + self.assertFalse(result.wasSuccessful()) + self.assertEqual([], result.errors) + self.assertEqual(1, len(result.failures)) + + def test_required_slice_contract_accepts_leading_indentation(self) -> None: + script = _ci_script() + invocation = ( + "run_required_gtests env WEBGPU_REQUIRE_DEVICE=1 " + "WEBGPU_TEST_HEAVY=1 \\\n" + ) + self.assertEqual(1, script.count(invocation)) + result = _required_slice_contract_result( + script.replace(invocation, " " + invocation) + ) + + self.assertTrue(result.wasSuccessful()) + self.assertEqual([], result.errors) + self.assertEqual([], result.failures) + + def test_dynamic_slice_export_is_fresh_heavy_and_fail_closed(self) -> None: + script = _ci_script() + recreate = ( + 'recreate_exact_directory "${DYNAMIC_SHAPE_DIR}" ' '"/tmp/dynamic_shape"' + ) + heavy_export = ( + '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" + ) + next_export = ( + '$PYTHON_EXECUTABLE -c "\n' + "from executorch.backends.webgpu.test.ops.test_sdpa import (" + ) + closure = script[script.index(recreate) : script.index(next_export)] + + self.assertIn(heavy_export, closure) + self.assertLess(closure.index(recreate), closure.index(heavy_export)) + for fixture in ( + "${DYNAMIC_SHAPE_DIR}/dyn_slice_2d.pte", + "${DYNAMIC_SHAPE_DIR}/slice_dual_store.pte", + "${DYNAMIC_SHAPE_DIR}/slice_dual_store.input.bin", + "${DYNAMIC_SHAPE_DIR}/slice_dual_store.out0.golden.bin", + "${DYNAMIC_SHAPE_DIR}/slice_dual_store.out1.golden.bin", + ): + requirement = f'require_file "{fixture}"' + self.assertEqual(script.count(requirement), 1, fixture) + self.assertIn(requirement, closure) + self.assertLess(closure.index(heavy_export), closure.index(requirement)) + + def test_dynamic_slice_export_cannot_reuse_stale_directory(self) -> None: + script = _ci_script() + with tempfile.TemporaryDirectory() as temporary: + root = pathlib.Path(temporary) + target = root / "dynamic_shape" + target.mkdir() + stale = target / "dyn_slice_2d.pte" + stale.write_text("stale") + + recreated = _run_recreate_exact_directory(script, target, target) + self.assertEqual(0, recreated.returncode, recreated.stderr) + self.assertTrue(target.is_dir()) + self.assertFalse(stale.exists()) + + protected = root / "protected" + protected.mkdir() + sentinel = protected / "sentinel" + sentinel.write_text("keep") + rejected = _run_recreate_exact_directory(script, protected, target) + self.assertNotEqual(0, rejected.returncode) + self.assertTrue(sentinel.is_file()) + + def test_runs_required_slice_regressions_fail_closed(self) -> None: + script = _ci_script() + helper = _bash_function(script, "run_required_gtests") + filter_value = ( + "DynamicShape.SliceCrosses2dDispatchBoundary:" + "DynamicShape.CatCrosses2dDispatchBoundary:" + "DynamicShape.SliceDualStoreWritesBothDestinations" + ) + + self.assertIn('run_with_required_device "$@"', helper) + self.assertIn("DynamicShape.SliceCrosses2dDispatchBoundary", helper) + self.assertIn("DynamicShape.CatCrosses2dDispatchBoundary", helper) + self.assertIn("DynamicShape.SliceDualStoreWritesBothDestinations", helper) + self.assertIn("[ PASSED ] 3 tests.", helper) + self.assertIn("grep -Eq '^\\[ SKIPPED \\]'", helper) + self.assertEqual( + ( + "env", + "WEBGPU_REQUIRE_DEVICE=1", + "WEBGPU_TEST_HEAVY=1", + "/contract bin/webgpu_dynamic_shape_test", + "/contract dynamic shape", + f"--gtest_filter={filter_value}", + ), + _required_gtest_invocation(script), + ) + self.assertLess( + script.index( + '"${BIN_DIR}/webgpu_dynamic_shape_test" "${DYNAMIC_SHAPE_DIR}"' + ), + script.index("run_required_gtests env WEBGPU_REQUIRE_DEVICE=1"), + ) + + def test_update_cache_dynamic_contract_is_fail_closed(self) -> None: + script = _ci_script() + recreate = 'recreate_exact_directory "${UPDATE_CACHE_DIR}" "/tmp/update_cache"' + export_dynamic = ( + "export_dynamic_update_cache('${UPDATE_CACHE_DIR}/dynamic.pte')" + ) + export_intermediate = ( + "export_intermediate_dynamic_update_cache(" + "'${UPDATE_CACHE_DIR}/dynamic_intermediate.pte')" + ) + state_run = '"${BIN_DIR}/webgpu_update_cache_state_test"' + required_run = ( + "run_with_required_device env WEBGPU_REQUIRE_DEVICE=1 \\\n" + ' WEBGPU_UPDATE_CACHE_DIR="${UPDATE_CACHE_DIR}" \\\n' + ' "${BIN_DIR}/webgpu_update_cache_test" "${UPDATE_CACHE_DIR}"' + ) + + for fragment in ( + recreate, + export_dynamic, + export_intermediate, + 'require_file "${UPDATE_CACHE_DIR}/dynamic.pte"', + 'require_file "${UPDATE_CACHE_DIR}/dynamic_intermediate.pte"', + state_run, + required_run, + ): + self.assertEqual(script.count(fragment), 1, fragment) + self.assertLess(script.index(recreate), script.index(export_dynamic)) + self.assertLess(script.index(export_dynamic), script.index(state_run)) + self.assertLess(script.index(state_run), script.index(required_run)) + + source = _UPDATE_CACHE_TEST.read_text() + self.assertIn("test/native/RequiredDevicePolicy.h>", source) + self.assertIn('std::getenv("WEBGPU_REQUIRE_DEVICE")', source) + self.assertIn("required_device_failure_exit_code", source) + self.assertIn('std::printf("WebGPU device acquired (native)\\n")', source) + + def test_required_gtest_helper_rejects_incomplete_output(self) -> None: + script = _ci_script() + first = "DynamicShape.SliceCrosses2dDispatchBoundary" + second = "DynamicShape.CatCrosses2dDispatchBoundary" + third = "DynamicShape.SliceDualStoreWritesBothDestinations" + accepted = "\n".join( + ( + "WebGPU device acquired (native)", + f"[ OK ] {first} (1 ms)", + f"[ OK ] {second} (2 ms)", + f"[ OK ] {third} (3 ms)", + "[ PASSED ] 3 tests.", + ) + ) + + passed = _run_required_gtests(script, accepted) + self.assertEqual(0, passed.returncode, passed.stderr) + rejected = ( + (accepted.replace("WebGPU device acquired (native)\n", ""), 0), + (accepted.replace(f"[ OK ] {second} (2 ms)\n", ""), 0), + (accepted.replace(f"[ OK ] {third} (3 ms)\n", ""), 0), + (accepted.replace("[ PASSED ] 3 tests.", "[ PASSED ] 2 tests."), 0), + (accepted + "\n[ SKIPPED ] DynamicShape.Unexpected (0 ms)", 0), + (accepted, 3), + ) + for output, status in rejected: + with self.subTest(output=output, status=status): + result = _run_required_gtests(script, output, status) + self.assertNotEqual(0, result.returncode, result.stdout) + + def test_dynamic_slice_main_fails_closed_when_device_is_required(self) -> None: + source = _DYNAMIC_SHAPE_TEST.read_text() + + self.assertIn( + "test/native/RequiredDevicePolicy.h>", + source, + ) + self.assertIn('std::getenv("WEBGPU_REQUIRE_DEVICE")', source) + self.assertIn("required_device_failure_exit_code", source) + self.assertIn('std::printf("WebGPU device acquired (native)\\n")', source) + + def test_slice_dispatch_grid_helper_owns_both_dimensions(self) -> None: + self.assertTrue(_SLICE_DISPATCH.is_file(), _SLICE_DISPATCH) + header = _SLICE_DISPATCH.read_text() + implementation = _SLICE_IMPL.read_text() + + self.assertIn("dispatch.workgroup_count_x = grid.x;", header) + self.assertIn("dispatch.workgroup_count_y = grid.y;", header) + self.assertEqual(implementation.count("set_slice_dispatch_grid("), 2) + self.assertNotIn("workgroup_count_y = wgc.y", implementation) + + def test_slice_correctness_does_not_require_profiling(self) -> None: + source = _DYNAMIC_SHAPE_TEST.read_text() + profile_available = source.split("bool slice_profile_available()", 1)[1].split( + "void expect_slice_profile", 1 + )[0] + boundary = source.split( + "TEST(DynamicShape, SliceCrosses2dDispatchBoundary)", 1 + )[1].split("TEST(DynamicShape, SliceDualStoreWritesBothDestinations)", 1)[0] + dual_store = source.split( + "TEST(DynamicShape, SliceDualStoreWritesBothDestinations)", 1 + )[1].split("TEST(DynamicShape, ExpandCopyRejectsDynamicShapesAtLoad)", 1)[0] + + self.assertNotIn("timestamp queries unavailable", boundary) + self.assertNotIn("timestamp queries unavailable", dual_store) + for predicate in ( + 'std::getenv("WEBGPU_TIMESTAMP_QUERY") != nullptr', + "context != nullptr", + "context->timestamp_supported", + "context->querypool != nullptr", + ): + self.assertIn(predicate, profile_available) + self.assertEqual(source.count("if (slice_profile_available())"), 2) + + def test_runs_codegen_pin_gate_before_fixture_exports(self) -> None: + script = _ci_script() + command = ( + "buck2 test " "fbcode//executorch/backends/webgpu/test:test_wgsl_codegen" + ) + + self.assertIn( + "test_wgsl_codegen", + _target_names((_BACKEND / "test/BUCK").read_text()), + ) + self.assertEqual(script.count(command), 1) + self.assertLess(script.index(command), script.index("# ── Exports")) + + def test_persistently_validates_wasm_names_without_claiming_products( + self, + ) -> None: + script = _ci_script() + invocation = ( + 'bash "${SCRIPT_DIR}/test_gemma4_wasm_factory_contract.sh" ' + "--validate-names" + ) + self.assertEqual(script.count(invocation), 1) + self.assertNotIn("--verify-product", script) + self.assertLess(script.index(invocation), script.index("# ── Exports")) + def test_builds_and_runs_every_fixed_target(self) -> None: - backend = pathlib.Path(__file__).parents[1] - cmake = (backend / "CMakeLists.txt").read_text() - script = (backend / "scripts/test_webgpu_native_ci.sh").read_text() + cmake = _CMAKE.read_text() + script = _ci_script() required = { "webgpu_native_test", "webgpu_dispatch_order_test", "webgpu_scratch_buffer_test", "webgpu_update_cache_test", + "webgpu_update_cache_state_test", "webgpu_index_test", "webgpu_dynamic_shape_test", "webgpu_dispatch_2d_test", @@ -36,6 +755,9 @@ def test_builds_and_runs_every_fixed_target(self) -> None: "webgpu_execution_options_test", "webgpu_output_suppression_test", "webgpu_op_test_util_test", + "webgpu_topk_test", + "webgpu_scatter_test", + "webgpu_q4gsw_m3_test", } self.assertEqual(set(_bash_array(script, "REQUIRED_TARGETS")), required) @@ -47,13 +769,11 @@ def test_builds_and_runs_every_fixed_target(self) -> None: script, ) for target in required: - self.assertIn(target, cmake) + self.assertTrue(_cmake_defines(cmake, target), target) self.assertIn(f'"${{BIN_DIR}}/{target}"', script) def test_requires_symint_and_suppression_fixtures(self) -> None: - script = ( - pathlib.Path(__file__).parents[1] / "scripts/test_webgpu_native_ci.sh" - ).read_text() + script = _ci_script() self.assertIn( "export_output_suppression_models('${OUTPUT_SUPPRESSION_DIR}')", script @@ -66,9 +786,7 @@ def test_requires_symint_and_suppression_fixtures(self) -> None: self.assertIn(f'require_file "{fixture}"', script) def test_requires_dynamic_rope_fixture(self) -> None: - script = ( - pathlib.Path(__file__).parents[1] / "scripts/test_webgpu_native_ci.sh" - ).read_text() + script = _ci_script() self.assertIn("export_rope_hf_dynamic('${ROPE_HF_DIR}')", script) self.assertIn('WEBGPU_TEST_ROPE_HF_DIR="${ROPE_HF_DIR}"', script) @@ -113,3 +831,116 @@ def test_cat_2d_regressions_are_heavy_and_fail_closed(self) -> None: 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) + + def test_exports_and_requires_topk_and_scatter_fixtures(self) -> None: + script = _ci_script() + + self.assertIn( + 'EAGLE_TOPK_EAGER_RECEIPT="${TOPK_AUTHORITY}" $PYTHON_EXECUTABLE ' + "-m unittest", + script, + ) + self.assertIn( + "executorch.backends.webgpu.test.ops.topk.test_topk.TestEagleTopKCpu" + ".test_eager_reference_is_repeatable", + script, + ) + self.assertIn( + "-m executorch.backends.webgpu.test.ops.topk.export_topk_artifacts", + script, + ) + self.assertIn('"${TOPK_DIR}" "${TOPK_AUTHORITY}"', script) + self.assertIn( + "-m executorch.backends.webgpu.test.ops.scatter.export_scatter_artifacts", + script, + ) + for fixture in ( + "${TOPK_AUTHORITY}", + "${TOPK_DIR}/cases.txt", + "${SCATTER_DIR}/cases.txt", + "${SCATTER_DIR}/base.bin", + ): + self.assertIn(f'require_file "{fixture}"', script) + + self.assertIn('"${BIN_DIR}/webgpu_topk_test" "${TOPK_DIR}"', script) + self.assertIn('"${BIN_DIR}/webgpu_scatter_test" "${SCATTER_DIR}"', script) + + def test_every_registered_source_exists(self) -> None: + missing: list[str] = [] + for build_file in _BUILD_FILES: + entries = _srcs_entries(build_file.read_text()) + self.assertNotEqual( + entries, [], f"{build_file} registers no srcs; parser drifted" + ) + for entry in entries: + if not (build_file.parent / entry).is_file(): + missing.append(f"{build_file}: {entry}") + self.assertEqual( + missing, [], f"registrations point at missing sources: {missing}" + ) + + def test_every_d10_test_source_is_registered(self) -> None: + registered = _registered_sources() + unregistered = [name for name in _D10_TEST_SOURCES if name not in registered] + self.assertEqual( + unregistered, + [], + f"no Buck srcs, CMake target, or CI invocation names: {unregistered}", + ) + + def test_committed_d10_test_sources_match_the_uncommitted_working_copy( + self, + ) -> None: + """Strict where source control can still identify D10; else says so.""" + reported = _reported_test_sources() + if reported is None: + self.assertFalse( + _sapling_is_usable(), + "`sl status` failed inside a Sapling working copy, so the " + "committed D10 test source list went unverified here", + ) + return + if _THIS_GATE not in reported: + self.assertNotEqual( + _sl_log_node(_THIS_GATE), + "", + "this gate is neither an uncommitted change nor a committed " + "file, so the committed D10 test source list went unverified", + ) + return + self.assertEqual( + reported, + list(_D10_TEST_SOURCES), + "the committed D10 test source list drifted from the working copy", + ) + + def test_mtp_command_contract_targets_are_defined(self) -> None: + undefined = _undefined_labels(_MTP_BUCK_TEST_TARGETS) + self.assertEqual(undefined, [], f"MTP labels are undefined: {undefined}") + + def test_plain_regression_command_contract_targets_are_defined(self) -> None: + undefined = _undefined_labels(_PLAIN_REGRESSION_TARGETS) + self.assertEqual( + undefined, [], f"plain-regression labels are undefined: {undefined}" + ) + + def test_documented_cmake_build_targets_are_defined(self) -> None: + cmake = _CMAKE.read_text() + missing = [ + name for name in _CMAKE_BUILD_TARGETS if not _cmake_defines(cmake, name) + ] + self.assertEqual( + missing, [], f"CMakeLists.txt defines no such target: {missing}" + ) + + def test_orphaned_test_targets_bzl_stays_a_pure_duplicate(self) -> None: + buck = (_BACKEND / "test/BUCK").read_text() + orphan = (_BACKEND / "test/targets.bzl").read_text() + + self.assertNotIn('load(":targets.bzl"', buck) + self.assertEqual( + _target_names(orphan) - _target_names(buck), + set(), + "test/BUCK does not load test/targets.bzl, so a target only the " + "latter defines would never be built", + ) diff --git a/backends/webgpu/test/test_wgsl_codegen.py b/backends/webgpu/test/test_wgsl_codegen.py index 244ab6eb67a..989672a4a37 100644 --- a/backends/webgpu/test/test_wgsl_codegen.py +++ b/backends/webgpu/test/test_wgsl_codegen.py @@ -249,11 +249,11 @@ def test_generated_output_manifest_digest(self) -> None: self.assertEqual(len(outputs), 149) self.assertEqual( digest.hexdigest(), - "8b2879a6ba11b57fd67aa961793ef9ff5142fdefaa7d9dcf41ab26276331f546", + "6854cdb9e33cdfe638edebe8ddc083cbf612e0a366870b275cd3b908b92555fc", ) self.assertEqual( hashlib.sha256(g.registry_path().read_bytes()).hexdigest(), - "1c26ac3f0671aeec5c648f78c9f2cbeb02b65f79478712a4f4de5c7e75446e8c", + "774ec59a0d1f17138b09090ea0bc76645f3d9fdd816fa988c7a4b9329fb77b43", ) def test_rope_hf_reconstructs_full_2d_grid_stride(self) -> None: @@ -1105,6 +1105,10 @@ def test_binary_family_roundtrip_byte_identical(self) -> None: 0, "63209ff70422a21fc340d9aadba0945bc259bba89bdf05db018a6507d01c7ae5", ), + "binary_sub_int32": ( + 0, + "134151da070a891e539f6ede5974310c623a9a7d379c90e69f91bec56ddc9b29", + ), "binary_minimum": ( 1, "929b7ba85936e3652baea9f4e5e7f049d232c7ae7a74814a536b4c2674897972", @@ -1143,6 +1147,21 @@ def test_binary_family_roundtrip_byte_identical(self) -> None: f"runtime/ops/binary_op/{name}_wgsl.h", ) + int32_params = variants["binary_sub_int32"] + self.assertEqual(int32_params["SCALAR_TYPE"], "i32") + self.assertEqual(int32_params["ALPHA_TYPE"], "i32") + self.assertEqual(int32_params["ALPHA_DEFAULT"], "1i") + self.assertEqual( + int32_params["OP_EXPR"], + "bitcast(bitcast(a) - bitcast(alpha) * bitcast(b))", + ) + self.assertFalse( + (g.BACKEND_ROOT / "runtime/ops/sub/binary_sub_int32.wgsl").exists() + ) + self.assertFalse( + (g.BACKEND_ROOT / "runtime/ops/sub/binary_sub_int32_wgsl.h").exists() + ) + def test_unary_template_roundtrip_byte_identical(self) -> None: unary_dir = g.BACKEND_ROOT / "runtime/ops/unary" template_path = unary_dir / "unary.wgsl" diff --git a/examples/models/gemma4/README.md b/examples/models/gemma4/README.md index cf132e1ebd5..d1df1f9afe9 100644 --- a/examples/models/gemma4/README.md +++ b/examples/models/gemma4/README.md @@ -447,6 +447,22 @@ validation, and the combined view remains pending cross-view GPU execution validation; source-bound bytes alone convey no correctness or performance claim. +Generate the MTP replay oracle from the production manifest and the two +content-verified checkpoint directories. The generator requires the production +manifest to be the same MTP receipt staged in the combined runtime envelope. + +```bash +PYTHONPATH="$FBSOURCE_ROOT/xplat" python \ + "$FBSOURCE_ROOT/xplat/executorch/examples/models/gemma4/tests/generate_mtp_spec_oracle.py" \ + --production-mtp-manifest "$MTP_MANIFEST" \ + --target-checkpoint /tmp/gemma4-e2b-it \ + --assistant-checkpoint /tmp/gemma4-e2b-assistant \ + --combined-runtime-root staged \ + --authority portable_eager \ + --contexts 128,511,512,513,514,1024,2048,4096,4097,8192 \ + --output gemma4-mtp-spec-oracle.json +``` + The browser adapter exports load, reset, prefill, decode, profiling, and unload entry points through the `gemma4_spec_browser` CMake target. A reset clears controller state and unloads and reloads `k2_round`. Unload destroys diff --git a/examples/models/gemma4/tests/generate_mtp_spec_oracle.py b/examples/models/gemma4/tests/generate_mtp_spec_oracle.py new file mode 100644 index 00000000000..46b23f171e4 --- /dev/null +++ b/examples/models/gemma4/tests/generate_mtp_spec_oracle.py @@ -0,0 +1,1163 @@ +# 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. + +# pyre-strict + +"""Portable/eager replay bound to target-only Gemma 4 prefill evidence. + +The replay shares the K=2 model definition with the exporter. The staged +target-only receipt independently covers the eager target prefill path, not the +shared Gemma model implementation. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import math + +from pathlib import Path +from typing import Any, Mapping, Sequence + +from executorch.examples.models.gemma4.target_prefill_contract import ( + canonical_json_bytes, + prompt_plan_sha256, + reviewed_producer_source_path, + TARGET_PREFILL_AUTHORITY, + TARGET_PREFILL_CONTEXTS, + TARGET_PREFILL_ENVELOPE_KIND, + TARGET_PREFILL_SCHEMA_VERSION, + validate_target_prefill_receipt as validate_target_prefill_v2_receipt, +) +from executorch.examples.models.gemma4.webgpu_artifact_manifest import ( + CHECKPOINT_ACQUISITION, + MTP_SOURCE_VERIFIED_PROVENANCE, + validate_combined_runtime_envelope, +) + +ORACLE_SCHEMA_VERSION = 2 +LEGACY_ORACLE_SCHEMA_VERSION = 1 +MANIFEST_SCHEMA_VERSION = 1 +RECEIPT_SCHEMA_VERSION = 1 +SUPPORTED_AUTHORITIES = ("portable_eager",) +TARGET_PREFILL_AUTHORITIES: tuple[str, ...] = (TARGET_PREFILL_AUTHORITY,) +CLOSURE_STATES = ("absent", "full") +TARGET_PREFILL_BINDING_STATES = ("bound", "legacy_unbound") + +K2_METHOD_NAME = "k2_round" +K2_DRAFT_COUNT = 2 +K2_GREEDY_COUNT = 3 +K2_MIN_START_POSITION = 2 +K2_VOCAB_SIZE = 262144 +K2_MAX_SEQ_LEN = 8960 +K2_MAX_INPUT_LEN = 512 +K2_PTD_COUNT = 3 +ORACLE_TOKEN_BUDGET = 32 + +K2_ROUND_ABI: dict[str, Any] = { + "buffer_mutation_count": 31, + "operator_counts": { + "aten.argmax.default": 3, + "aten.scatter.src": 2, + "aten.topk.default": 2, + "llama.custom_sdpa.default": 43, + }, + "seed_mutation_count": 1, + "user_inputs": ["input_ids", "input_pos", "is_round", "donor_length"], + "user_outputs": [ + {"dtype": "int64", "name": "candidates", "shape": [1, 2]}, + {"dtype": "int64", "name": "target_greedy", "shape": [1, 3]}, + {"dtype": "int64", "name": "output_matches", "shape": [1]}, + {"dtype": "int64", "name": "output_bonus", "shape": [1, 1]}, + {"dtype": "float32", "name": "state_probe", "shape": [1, 1]}, + ], +} + +CHECKPOINT_ROLES = ("assistant", "target") + +TARGET_PREFILL_WITNESS_KEYS = ( + "layer0_av_sha256", + "layer0_qk_sha256", + "logits_sha256", + "prefill_token", +) + +ORACLE_CONTEXT_KEYS = ( + "accepted_prefix", + "bonus_accounting", + "decoded_text", + "kv_witnesses", + "reset_replay", + "rounds", + "selected_logits", + "stop_handling", + "target_prefill", + "useful_tokens", +) + +_ORACLE_TOP_LEVEL_KEYS = { + "abi", + "authority", + "closure_state", + "contexts", + "method", + "mtp_manifest_sha256", + "production_binding", + "records", + "replay_independence", + "schema_version", + "stop_tokens", + "target_prefill_authority", + "target_prefill_oracle_sha256", + "token_budget", +} +_PRODUCTION_BINDING_KEYS = { + "checkpoint_acquisition", + "combined_runtime_sha256", + "mtp_manifest_sha256", + "mtp_provenance", + "producer", + "run", + "target_prefill_receipt_sha256", +} +_RAW_ROUND_KEYS = { + "bonus", + "candidates", + "match_count", + "state_probe", + "target_greedy", +} +_DECISION_KEYS = { + "accepted_drafts", + "committed", + "discarded", + "next_position", + "next_seed", + "selected", + "stop_token", + "stopped", + "valid", +} +_ROUND_KEYS = _RAW_ROUND_KEYS | _DECISION_KEYS | {"kv_witness"} + + +class OracleError(Exception): + """Fail-closed rejection raised before any oracle bytes are written.""" + + +def _is_exact_int(value: object) -> bool: + return isinstance(value, int) and not isinstance(value, bool) + + +def _is_hex_digest(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, Any]: + if not isinstance(value, dict): + raise OracleError(f"{label} must be an object") + return value + + +def _sequence(value: object, label: str) -> Sequence[Any]: + if not isinstance(value, list): + raise OracleError(f"{label} must be a list") + return value + + +def _require_exact_keys( + value: Mapping[str, Any], expected: set[str], label: str +) -> None: + if set(value) != expected: + raise OracleError(f"{label} has an unexpected key set") + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as source: + for block in iter(lambda: source.read(1024 * 1024), b""): + digest.update(block) + return digest.hexdigest() + + +def load_json_object(path: Path, label: str) -> dict[str, Any]: + try: + text = path.read_text(encoding="utf-8") + except OSError as error: + raise OracleError(f"unreadable {label}: {path}") from error + try: + document = json.loads(text) + except json.JSONDecodeError as error: + raise OracleError(f"malformed {label}: {path}") from error + if not isinstance(document, dict): + raise OracleError(f"{label} must be a JSON object: {path}") + return document + + +def require_schema_version( + document: Mapping[str, Any], label: str, expected: int +) -> None: + version = document.get("schema_version") + if not _is_exact_int(version) or version != expected: + raise OracleError(f"{label} schema_version must be {expected}, got {version!r}") + + +def parse_contexts(value: str, max_context_length: int) -> list[int]: + fields = [field.strip() for field in value.split(",")] + if not fields or any(not field for field in fields): + raise OracleError(f"--contexts must be a non-empty integer list: {value!r}") + contexts: list[int] = [] + for field in fields: + if not field.isdigit(): + raise OracleError(f"--contexts entry is not a decimal integer: {field!r}") + context = int(field) + if context <= 0 or context > max_context_length: + raise OracleError( + f"--contexts entry {context} is outside 1..{max_context_length}" + ) + if context in contexts: + raise OracleError(f"--contexts entry is duplicated: {context}") + contexts.append(context) + return contexts + + +def _validate_checkpoints(manifest: Mapping[str, Any], root: Path) -> dict[str, Path]: + checkpoints = manifest.get("checkpoints") + if not isinstance(checkpoints, dict): + raise OracleError("MTP manifest is missing the checkpoints binding") + if tuple(sorted(checkpoints)) != CHECKPOINT_ROLES: + raise OracleError( + f"MTP manifest checkpoints must name exactly {list(CHECKPOINT_ROLES)}" + ) + resolved: dict[str, Path] = {} + for role in CHECKPOINT_ROLES: + value = checkpoints[role] + if not isinstance(value, str) or not value: + raise OracleError(f"MTP manifest {role} checkpoint must be a path") + path = Path(value) + candidate = path if path.is_absolute() else root / path + if not candidate.is_dir(): + raise OracleError(f"MTP manifest {role} checkpoint is not a directory") + resolved[role] = candidate + return resolved + + +def _validate_artifact_roles(manifest: Mapping[str, Any]) -> None: + artifacts = manifest.get("artifacts") + ptd_order = manifest.get("ptd_order") + if not isinstance(artifacts, list) or not isinstance(ptd_order, list): + raise OracleError("MTP manifest artifacts/PTD order must be lists") + roles = [entry.get("role") for entry in artifacts if isinstance(entry, dict)] + if len(roles) != len(artifacts): + raise OracleError("MTP manifest artifact entries must be objects") + if roles.count("pte") != 1: + raise OracleError("MTP manifest requires exactly one PTE artifact") + if roles.count("ptd") != K2_PTD_COUNT or len(ptd_order) != K2_PTD_COUNT: + raise OracleError(f"MTP manifest requires exactly {K2_PTD_COUNT} ordered PTDs") + + +def load_stop_tokens(target_checkpoint: Path) -> list[int]: + document = load_json_object( + target_checkpoint / "generation_config.json", "generation config" + ) + value = document.get("eos_token_id") + tokens = value if isinstance(value, list) else [value] + if not tokens or any( + not _is_exact_int(token) or token < 0 or token >= K2_VOCAB_SIZE + for token in tokens + ): + raise OracleError(f"generation config eos_token_id is invalid: {value!r}") + if len(set(tokens)) != len(tokens): + raise OracleError("generation config eos_token_id contains duplicates") + return list(tokens) + + +def validate_mtp_manifest( + manifest: Mapping[str, Any], manifest_path: Path +) -> dict[str, Any]: + require_schema_version(manifest, "MTP manifest", MANIFEST_SCHEMA_VERSION) + if manifest.get("method") != K2_METHOD_NAME: + raise OracleError( + f"MTP manifest method must be {K2_METHOD_NAME!r}, " + f"got {manifest.get('method')!r}" + ) + if manifest.get("abi") != K2_ROUND_ABI: + raise OracleError("MTP manifest does not match the D8 K=2 graph/ABI contract") + _validate_artifact_roles(manifest) + max_context_length = manifest.get("max_context_length") + if ( + not _is_exact_int(max_context_length) + or max_context_length <= 0 + or max_context_length > K2_MAX_SEQ_LEN + ): + raise OracleError( + f"MTP manifest max_context_length must be 1..{K2_MAX_SEQ_LEN}, " + f"got {max_context_length!r}" + ) + checkpoints = _validate_checkpoints(manifest, manifest_path.parent) + return { + "checkpoints": checkpoints, + "max_context_length": max_context_length, + "stop_tokens": load_stop_tokens(checkpoints["target"]), + } + + +def validate_target_prefill_receipt( + receipt: Mapping[str, Any], authority: str, contexts: Sequence[int] +) -> dict[str, Mapping[str, Any]]: + require_schema_version(receipt, "target-prefill receipt", RECEIPT_SCHEMA_VERSION) + if receipt.get("authority") != authority: + raise OracleError( + f"target-prefill receipt authority is {receipt.get('authority')!r}, " + f"not {authority!r}" + ) + entries = receipt.get("contexts") + if not isinstance(entries, dict): + raise OracleError("target-prefill receipt contexts must be an object") + witnesses: dict[str, Mapping[str, Any]] = {} + for context in contexts: + key = str(context) + entry = entries.get(key) + if not isinstance(entry, dict): + raise OracleError(f"target-prefill receipt lacks context {context}") + if tuple(sorted(entry)) != TARGET_PREFILL_WITNESS_KEYS: + raise OracleError( + f"target-prefill witness {context} must name exactly " + f"{list(TARGET_PREFILL_WITNESS_KEYS)}" + ) + token = entry["prefill_token"] + if not _is_exact_int(token) or token < 0 or token >= K2_VOCAB_SIZE: + raise OracleError(f"target-prefill witness {context} token is invalid") + for name in TARGET_PREFILL_WITNESS_KEYS: + if name != "prefill_token" and not _is_hex_digest(entry[name], 64): + raise OracleError( + f"target-prefill witness {context} {name} is not a digest" + ) + witnesses[key] = entry + return witnesses + + +def _combined_receipt_path(root: Path, envelope: Mapping[str, Any], role: str) -> Path: + receipts = envelope.get("receipts") + if not isinstance(receipts, dict): + raise OracleError("combined runtime receipts must be an object") + identity = receipts.get(role) + if not isinstance(identity, dict) or not isinstance(identity.get("path"), str): + raise OracleError(f"combined runtime lacks the {role} receipt identity") + resolved_root = root.resolve(strict=True) + path = (root / identity["path"]).resolve(strict=True) + try: + path.relative_to(resolved_root) + except ValueError as error: + raise OracleError( + f"combined runtime {role} receipt escapes its root" + ) from error + if not path.is_file(): + raise OracleError(f"combined runtime {role} receipt is not a regular file") + return path + + +def _validate_production_target_prefill_receipt( + receipt: Mapping[str, Any], contexts: Sequence[int] +) -> dict[str, Mapping[str, Any]]: + if receipt.get("schema_version") != TARGET_PREFILL_SCHEMA_VERSION: + raise OracleError("production target-prefill receipt requires schema version 2") + if receipt.get("envelope_kind") != TARGET_PREFILL_ENVELOPE_KIND: + raise OracleError("production target-prefill envelope kind mismatch") + if receipt.get("authority") not in TARGET_PREFILL_AUTHORITIES: + raise OracleError("production target-prefill authority mismatch") + if tuple(contexts) != TARGET_PREFILL_CONTEXTS: + raise OracleError("production binding requires the exact ten contexts") + entries = receipt.get("contexts") + if not isinstance(entries, dict) or set(entries) != { + str(context) for context in TARGET_PREFILL_CONTEXTS + }: + raise OracleError("production receipt requires the exact ten contexts") + witnesses: dict[str, Mapping[str, Any]] = {} + for context in TARGET_PREFILL_CONTEXTS: + entry = entries.get(str(context)) + if not isinstance(entry, dict): + raise OracleError( + f"production target-prefill receipt lacks context {context}" + ) + raw = entry.get("prefill_token_raw") + post = entry.get("prefill_token_post_softcap") + if not _is_exact_int(raw) or raw < 0 or raw >= K2_VOCAB_SIZE or post != raw: + raise OracleError( + f"target-prefill context {context} raw/post-softcap token mismatch" + ) + if entry.get("prompt_plan_sha256") != prompt_plan_sha256(context): + raise OracleError(f"target-prefill context {context} prompt plan mismatch") + witnesses[str(context)] = entry + return witnesses + + +def validate_k2_abi_edge_census(edge_census: Mapping[str, int]) -> None: + expected = { + "custom_scatter": K2_ROUND_ABI["operator_counts"]["aten.scatter.src"], + "gemma_sdpa": K2_ROUND_ABI["operator_counts"]["llama.custom_sdpa.default"], + "topk": K2_ROUND_ABI["operator_counts"]["aten.topk.default"], + } + if any(edge_census.get(name) != count for name, count in expected.items()): + raise OracleError("production edge census drifted from the K=2 ABI") + + +def require_prefill_token_match( + context: int, actual_token: int, witness: Mapping[str, Any] +) -> None: + expected = witness.get("prefill_token_raw", witness.get("prefill_token")) + if actual_token != expected: + raise OracleError( + f"context {context} prefill token disagrees with the target-only receipt" + ) + + +def reconcile_k2_round( + candidates: Sequence[int], + target_greedy: Sequence[int], + match_count: int, + bonus: int, + state_probe: float, + start_position: int, + token_budget: int, + stop_tokens: Sequence[int], + vocab_size: int = K2_VOCAB_SIZE, +) -> dict[str, Any]: + """Portable mirror of `reconcile_gemma4_k2` in `gemma4_spec_runner.h`.""" + decision: dict[str, Any] = { + "accepted_drafts": 0, + "committed": [], + "discarded": [], + "next_position": -1, + "next_seed": -1, + "selected": [], + "stop_token": -1, + "stopped": False, + "valid": False, + } + if ( + len(candidates) != K2_DRAFT_COUNT + or len(target_greedy) != K2_GREEDY_COUNT + or start_position < K2_MIN_START_POSITION + or token_budget <= 0 + or vocab_size <= 0 + or match_count < 0 + or match_count > K2_DRAFT_COUNT + or not math.isfinite(state_probe) + ): + return decision + if any( + token < 0 or token >= vocab_size + for token in list(candidates) + list(target_greedy) + ): + return decision + + expected_matches = 0 + if candidates[0] == target_greedy[0]: + expected_matches = 2 if candidates[1] == target_greedy[1] else 1 + if ( + match_count != expected_matches + or bonus < 0 + or bonus >= vocab_size + or bonus != target_greedy[match_count] + ): + return decision + + selected = list(candidates[:match_count]) + [bonus] + committed: list[int] = [] + discarded: list[int] = [] + for index, token in enumerate(selected): + if token in stop_tokens: + decision["stopped"] = True + decision["stop_token"] = token + discarded.extend(selected[index + 1 :]) + break + if len(committed) == token_budget: + discarded.extend(selected[index:]) + break + committed.append(token) + + decision.update( + { + "accepted_drafts": match_count, + "committed": committed, + "discarded": discarded, + "next_position": start_position + match_count + 1, + "next_seed": bonus, + "selected": selected, + "valid": True, + } + ) + return decision + + +def _bind_combined_runtime( + binding: dict[str, Any], + combined_runtime_root: Path, + *, + expected_mtp_sha256: str | None = None, + expected_checkpoint_acquisition: Mapping[str, Any] | None = None, +) -> dict[str, Any]: + envelope_path = combined_runtime_root / "gemma4_webgpu_combined_runtime.json" + envelope = load_json_object(envelope_path, "combined runtime envelope") + try: + validate_combined_runtime_envelope(combined_runtime_root, envelope) + except ValueError as error: + raise OracleError(str(error)) from error + if envelope.get("schema_version") != 3: + raise OracleError("production binding requires combined runtime schema 3") + + mtp_receipt_path = _combined_receipt_path(combined_runtime_root, envelope, "mtp") + target_prefill_path = _combined_receipt_path( + combined_runtime_root, envelope, "target_prefill" + ) + mtp_sha256 = _sha256(mtp_receipt_path) + if expected_mtp_sha256 is not None and mtp_sha256 != expected_mtp_sha256: + raise OracleError("production MTP manifest is not the staged MTP receipt") + production_mtp = load_json_object(mtp_receipt_path, "production MTP manifest") + if production_mtp.get("provenance") != MTP_SOURCE_VERIFIED_PROVENANCE: + raise OracleError("production MTP manifest is not source verified") + target_receipt = load_json_object( + target_prefill_path, "production target-prefill receipt" + ) + if ( + expected_checkpoint_acquisition is not None + and target_receipt.get("checkpoint_acquisition") + != expected_checkpoint_acquisition + ): + raise OracleError("target-prefill checkpoint acquisition mismatch") + binding.update( + { + "closure_state": "full", + "mtp_manifest_sha256": mtp_sha256, + "production_binding": { + "checkpoint_acquisition": target_receipt.get( + "checkpoint_acquisition" + ), + "combined_runtime_sha256": _sha256(envelope_path), + "mtp_manifest_sha256": mtp_sha256, + "mtp_provenance": production_mtp["provenance"], + "producer": target_receipt.get("producer"), + "run": target_receipt.get("run"), + "target_prefill_receipt_sha256": _sha256(target_prefill_path), + }, + "target_prefill": _validate_production_target_prefill_receipt( + target_receipt, binding["contexts"] + ), + "target_prefill_authority": "bound", + "target_prefill_oracle_sha256": _sha256(target_prefill_path), + } + ) + return binding + + +def build_production_oracle_binding( + production_mtp_manifest: Path, + target_checkpoint: Path, + assistant_checkpoint: Path, + authority: str, + contexts: str, + output: Path, + *, + combined_runtime_root: Path, +) -> dict[str, Any]: + if authority not in SUPPORTED_AUTHORITIES: + raise OracleError( + f"unknown --authority {authority!r}; expected one of " + f"{list(SUPPORTED_AUTHORITIES)}" + ) + if output.exists() or output.is_symlink(): + raise OracleError(f"refusing to overwrite existing artifact: {output}") + manifest = load_json_object(production_mtp_manifest, "production MTP manifest") + if manifest.get("provenance") != MTP_SOURCE_VERIFIED_PROVENANCE: + raise OracleError("production MTP manifest is not source verified") + export = _mapping(manifest.get("export"), "production MTP export") + if export.get("methods") != [K2_METHOD_NAME]: + raise OracleError("production MTP manifest does not expose k2_round") + max_context_length = export.get("max_seq_len") + if ( + not _is_exact_int(max_context_length) + or max_context_length <= 0 + or max_context_length > K2_MAX_SEQ_LEN + ): + raise OracleError("production MTP max_seq_len is invalid") + _validate_artifact_roles(manifest) + acquisition = _mapping( + manifest.get("acquisition"), "production MTP acquisition" + ) + if acquisition.get("target") != CHECKPOINT_ACQUISITION: + raise OracleError("production MTP target acquisition mismatch") + checkpoints = _validate_checkpoints( + { + "checkpoints": { + "assistant": str(assistant_checkpoint), + "target": str(target_checkpoint), + } + }, + production_mtp_manifest.parent, + ) + resolved = parse_contexts(contexts, max_context_length) + mtp_sha256 = _sha256(production_mtp_manifest) + binding: dict[str, Any] = { + "authority": authority, + "checkpoints": checkpoints, + "contexts": resolved, + "closure_state": "absent", + "max_context_length": max_context_length, + "mtp_manifest_sha256": mtp_sha256, + "production_binding": None, + "stop_tokens": load_stop_tokens(checkpoints["target"]), + "target_prefill_authority": "legacy_unbound", + "token_budget": ORACLE_TOKEN_BUDGET, + } + return _bind_combined_runtime( + binding, + combined_runtime_root, + expected_mtp_sha256=mtp_sha256, + expected_checkpoint_acquisition=CHECKPOINT_ACQUISITION, + ) + + +def build_oracle_binding( + mtp_manifest: Path, + target_prefill_oracle: Path | None, + authority: str, + contexts: str, + output: Path, + *, + combined_runtime_root: Path | None = None, +) -> dict[str, Any]: + if authority not in SUPPORTED_AUTHORITIES: + raise OracleError( + f"unknown --authority {authority!r}; expected one of " + f"{list(SUPPORTED_AUTHORITIES)}" + ) + if output.exists() or output.is_symlink(): + raise OracleError(f"refusing to overwrite existing artifact: {output}") + manifest = load_json_object(mtp_manifest, "MTP manifest") + bound = validate_mtp_manifest(manifest, mtp_manifest) + resolved = parse_contexts(contexts, bound["max_context_length"]) + binding: dict[str, Any] = { + "authority": authority, + "checkpoints": bound["checkpoints"], + "contexts": resolved, + "closure_state": "absent", + "max_context_length": bound["max_context_length"], + "mtp_manifest_sha256": _sha256(mtp_manifest), + "production_binding": None, + "stop_tokens": bound["stop_tokens"], + "target_prefill_authority": "legacy_unbound", + "token_budget": ORACLE_TOKEN_BUDGET, + } + if combined_runtime_root is None: + if target_prefill_oracle is None: + raise OracleError("legacy binding requires a target-prefill receipt") + receipt = load_json_object(target_prefill_oracle, "target-prefill receipt") + binding["target_prefill"] = validate_target_prefill_receipt( + receipt, authority, resolved + ) + binding["target_prefill_oracle_sha256"] = _sha256(target_prefill_oracle) + return binding + + if target_prefill_oracle is not None: + raise OracleError( + "production binding derives the target-prefill receipt from the " + "combined runtime root" + ) + return _bind_combined_runtime(binding, combined_runtime_root) + + +def assemble_oracle_document( + binding: Mapping[str, Any], records: Mapping[str, Mapping[str, Any]] +) -> dict[str, Any]: + expected = [str(context) for context in binding["contexts"]] + if sorted(records) != sorted(expected): + raise OracleError("authority records do not cover the requested contexts") + for key in expected: + record = records[key] + if tuple(sorted(record)) != ORACLE_CONTEXT_KEYS: + raise OracleError(f"authority record {key} has an unexpected key set") + if record["target_prefill"] != binding["target_prefill"][key]: + raise OracleError(f"authority record {key} reinterprets the D6 receipt") + production = binding.get("closure_state") == "full" + if production and tuple(binding["contexts"]) != TARGET_PREFILL_CONTEXTS: + raise OracleError("production oracle requires the exact ten contexts") + document: dict[str, Any] = { + "abi": K2_ROUND_ABI, + "authority": binding["authority"], + "closure_state": binding.get("closure_state", "absent"), + "contexts": list(binding["contexts"]), + "method": K2_METHOD_NAME, + "mtp_manifest_sha256": binding["mtp_manifest_sha256"], + "records": {key: dict(records[key]) for key in expected}, + "replay_independence": "eager_vs_lowered_only", + "schema_version": ( + ORACLE_SCHEMA_VERSION if production else LEGACY_ORACLE_SCHEMA_VERSION + ), + "stop_tokens": list(binding["stop_tokens"]), + "target_prefill_authority": binding.get( + "target_prefill_authority", "legacy_unbound" + ), + "target_prefill_oracle_sha256": binding["target_prefill_oracle_sha256"], + "token_budget": binding["token_budget"], + } + if production: + document["production_binding"] = binding["production_binding"] + return document + + +def _validate_production_binding( + document: Mapping[str, Any], records: Mapping[str, Any] +) -> None: + binding = _mapping(document.get("production_binding"), "production binding") + _require_exact_keys(binding, _PRODUCTION_BINDING_KEYS, "production binding") + for key in ( + "combined_runtime_sha256", + "mtp_manifest_sha256", + "target_prefill_receipt_sha256", + ): + if not _is_hex_digest(binding.get(key), 64): + raise OracleError(f"production binding {key} is invalid") + if binding.get("mtp_manifest_sha256") != document.get("mtp_manifest_sha256"): + raise OracleError("production MTP manifest identity is inconsistent") + if binding.get("target_prefill_receipt_sha256") != document.get( + "target_prefill_oracle_sha256" + ): + raise OracleError("target-prefill receipt identity is inconsistent") + if binding.get("mtp_provenance") != MTP_SOURCE_VERIFIED_PROVENANCE: + raise OracleError("production binding MTP provenance is not source verified") + if binding.get("checkpoint_acquisition") != CHECKPOINT_ACQUISITION: + raise OracleError("production binding checkpoint acquisition mismatch") + + producer = _mapping(binding.get("producer"), "target-prefill producer") + runtime_source_identity = _mapping( + producer.get("runtime_source_receipt"), + "target-prefill runtime source identity", + ) + producer_path = reviewed_producer_source_path() + receipt = { + "authority": TARGET_PREFILL_AUTHORITY, + "checkpoint_acquisition": binding["checkpoint_acquisition"], + "contexts": { + key: _mapping(record, f"oracle record {key}").get("target_prefill") + for key, record in records.items() + }, + "envelope_kind": TARGET_PREFILL_ENVELOPE_KIND, + "producer": producer, + "run": binding.get("run"), + "schema_version": TARGET_PREFILL_SCHEMA_VERSION, + } + receipt_sha256 = hashlib.sha256(canonical_json_bytes(receipt)).hexdigest() + if receipt_sha256 != binding.get("target_prefill_receipt_sha256"): + raise OracleError("target-prefill receipt content does not match its identity") + try: + validate_target_prefill_v2_receipt( + receipt, + expected_checkpoint_acquisition=CHECKPOINT_ACQUISITION, + expected_producer_path=producer_path, + expected_producer_sha256=str(producer.get("source_sha256")), + expected_runtime_source_identity=runtime_source_identity, + expected_fbsource_commit=str(producer.get("fbsource_commit")), + ) + except ValueError as error: + raise OracleError(str(error)) from error + + +def _int_sequence(value: object, label: str) -> list[int]: + sequence = _sequence(value, label) + if any(not _is_exact_int(item) for item in sequence): + raise OracleError(f"{label} must contain only integers") + return list(sequence) + + +def _validate_oracle_record( + record: object, + *, + context: int, + stop_tokens: Sequence[int], + token_budget: int, +) -> None: + value = _mapping(record, f"oracle record {context}") + _require_exact_keys(value, set(ORACLE_CONTEXT_KEYS), f"oracle record {context}") + target_prefill = _mapping( + value.get("target_prefill"), f"oracle target-prefill witness {context}" + ) + prefill_token = target_prefill.get("prefill_token_raw") + if not _is_exact_int(prefill_token): + raise OracleError(f"oracle target-prefill token {context} is invalid") + useful_tokens = _int_sequence(value.get("useful_tokens"), "oracle useful tokens") + if not useful_tokens or useful_tokens[0] != prefill_token: + raise OracleError("oracle useful tokens do not start with the prefill token") + if len(useful_tokens) > token_budget: + raise OracleError("oracle useful tokens exceed the token budget") + + rounds = _sequence(value.get("rounds"), "oracle rounds") + if not rounds: + raise OracleError("production oracle requires at least one K=2 round") + accepted_prefix: list[int] = [] + bonus_accounting: list[int] = [] + kv_witnesses: list[str] = [] + selected: list[list[int]] = [] + discarded: list[list[int]] = [] + reconstructed = [prefill_token] + position = context + reset_replay: dict[str, Any] | None = None + stop_token: int | None = None + for index, round_value in enumerate(rounds): + round_record = _mapping(round_value, f"oracle round {context}:{index}") + _require_exact_keys( + round_record, _ROUND_KEYS, f"oracle round {context}:{index}" + ) + candidates = _int_sequence( + round_record.get("candidates"), "oracle round candidates" + ) + target_greedy = _int_sequence( + round_record.get("target_greedy"), "oracle round target_greedy" + ) + bonus = round_record.get("bonus") + match_count = round_record.get("match_count") + state_probe = round_record.get("state_probe") + if ( + not _is_exact_int(bonus) + or not _is_exact_int(match_count) + or not isinstance(state_probe, (int, float)) + or isinstance(state_probe, bool) + or not math.isfinite(float(state_probe)) + ): + raise OracleError("oracle round scalar fields are invalid") + remaining = token_budget - len(reconstructed) + if remaining <= 0: + raise OracleError("oracle carries rounds after exhausting its token budget") + decision = reconcile_k2_round( + candidates, + target_greedy, + match_count, + bonus, + float(state_probe), + position, + remaining, + stop_tokens, + ) + if not decision["valid"] or any( + round_record.get(key) != decision[key] for key in _DECISION_KEYS + ): + raise OracleError("oracle round disagrees with K=2 reconciliation") + kv_witness = round_record.get("kv_witness") + if not _is_hex_digest(kv_witness, 64): + raise OracleError("oracle round KV witness is invalid") + raw = {key: round_record[key] for key in _RAW_ROUND_KEYS} + if reset_replay is None: + reset_replay = raw + accepted_prefix.append(decision["accepted_drafts"]) + bonus_accounting.append(decision["next_seed"]) + kv_witnesses.append(kv_witness) + selected.append(decision["selected"]) + discarded.append(decision["discarded"]) + reconstructed.extend(decision["committed"]) + position = decision["next_position"] + if decision["stopped"]: + stop_token = decision["stop_token"] + if index != len(rounds) - 1: + raise OracleError("oracle carries rounds after a stop token") + + if value.get("accepted_prefix") != accepted_prefix: + raise OracleError("oracle accepted-prefix summary is inconsistent") + if value.get("bonus_accounting") != bonus_accounting: + raise OracleError("oracle bonus summary is inconsistent") + if value.get("kv_witnesses") != kv_witnesses: + raise OracleError("oracle KV-witness summary is inconsistent") + if value.get("selected_logits") != selected: + raise OracleError("oracle selected-token summary is inconsistent") + if value.get("reset_replay") != reset_replay: + raise OracleError("oracle reset replay is inconsistent") + if useful_tokens != reconstructed: + raise OracleError("oracle useful-token summary is inconsistent") + stop_handling = _mapping(value.get("stop_handling"), "oracle stop handling") + _require_exact_keys(stop_handling, {"discarded", "stop_token"}, "stop handling") + if stop_handling != {"discarded": discarded, "stop_token": stop_token}: + raise OracleError("oracle stop handling is inconsistent") + if not isinstance(value.get("decoded_text"), str): + raise OracleError("oracle decoded text must be a string") + + +def _validate_production_oracle_document(document: Mapping[str, Any]) -> None: + _require_exact_keys(document, _ORACLE_TOP_LEVEL_KEYS, "production oracle") + if document.get("schema_version") != ORACLE_SCHEMA_VERSION: + raise OracleError("production oracle schema version mismatch") + if document.get("method") != K2_METHOD_NAME or document.get("abi") != K2_ROUND_ABI: + raise OracleError("production oracle K=2 method/ABI mismatch") + if document.get("authority") not in SUPPORTED_AUTHORITIES: + raise OracleError("production oracle replay authority mismatch") + if document.get("replay_independence") != "eager_vs_lowered_only": + raise OracleError("production oracle independence claim mismatch") + if document.get("token_budget") != ORACLE_TOKEN_BUDGET: + raise OracleError("production oracle token budget mismatch") + for key in ("mtp_manifest_sha256", "target_prefill_oracle_sha256"): + if not _is_hex_digest(document.get(key), 64): + raise OracleError(f"production oracle {key} is invalid") + contexts = _int_sequence(document.get("contexts"), "oracle contexts") + if tuple(contexts) != TARGET_PREFILL_CONTEXTS: + raise OracleError("production oracle requires the exact ten contexts") + stop_tokens = _int_sequence(document.get("stop_tokens"), "oracle stop tokens") + if ( + not stop_tokens + or len(stop_tokens) != len(set(stop_tokens)) + or any(token < 0 or token >= K2_VOCAB_SIZE for token in stop_tokens) + ): + raise OracleError("production oracle stop tokens are invalid") + records = _mapping(document.get("records"), "oracle records") + if set(records) != {str(context) for context in TARGET_PREFILL_CONTEXTS}: + raise OracleError("production oracle records do not cover all contexts") + _validate_production_binding(document, records) + for context in TARGET_PREFILL_CONTEXTS: + _validate_oracle_record( + records[str(context)], + context=context, + stop_tokens=stop_tokens, + token_budget=ORACLE_TOKEN_BUDGET, + ) + + +def production_oracle_is_acceptable(document: Mapping[str, Any]) -> bool: + closure_state = document.get("closure_state") + if closure_state not in CLOSURE_STATES: + raise OracleError(f"unknown oracle closure_state: {closure_state!r}") + target_authority = document.get("target_prefill_authority") + if target_authority not in TARGET_PREFILL_BINDING_STATES: + raise OracleError(f"unknown target_prefill_authority: {target_authority!r}") + if ( + document.get("schema_version") != ORACLE_SCHEMA_VERSION + or closure_state != "full" + or target_authority != "bound" + or tuple(document.get("contexts", ())) != TARGET_PREFILL_CONTEXTS + or not isinstance(document.get("production_binding"), dict) + ): + return False + try: + _validate_production_oracle_document(document) + except (KeyError, OracleError, TypeError, ValueError): + return False + return True + + +def _round_inputs( + torch: Any, seed_token: int, start_position: int +) -> tuple[Any, Any, Any, Any]: + return ( + torch.tensor([[seed_token, 0, 0]], dtype=torch.long), + torch.arange(start_position, start_position + 3, dtype=torch.long), + torch.tensor([1], dtype=torch.long), + torch.tensor([[start_position]], dtype=torch.long), + ) + + +def _kv_witness(torch: Any, module: Any) -> str: + digest = hashlib.sha256() + for name, buffer in sorted(module.state_dict().items()): + if name.endswith("k_cache") or name.endswith("v_cache"): + digest.update(name.encode("utf-8")) + digest.update(buffer.detach().to(torch.float32).cpu().numpy().tobytes()) + return digest.hexdigest() + + +def _read_round(module: Any, inputs: tuple[Any, ...]) -> dict[str, Any]: + candidates, target_greedy, matches, bonus, state_probe = module(*inputs) + return { + "bonus": int(bonus.reshape(-1)[0].item()), + "candidates": [int(value) for value in candidates.reshape(-1).tolist()], + "match_count": int(matches.reshape(-1)[0].item()), + "state_probe": float(state_probe.reshape(-1)[0].item()), + "target_greedy": [int(value) for value in target_greedy.reshape(-1).tolist()], + } + + +def _prefill(module: Any, torch: Any, prompt: Sequence[int]) -> dict[str, Any]: + output: dict[str, Any] = {} + start = 0 + while start < len(prompt): + count = min(K2_MAX_INPUT_LEN, len(prompt) - start) + output = _read_round( + module, + ( + torch.tensor([list(prompt[start : start + count])], dtype=torch.long), + torch.arange(start, start + count, dtype=torch.long), + torch.tensor([0], dtype=torch.long), + torch.tensor([[2 if start == 0 else start]], dtype=torch.long), + ), + ) + start += count + return output + + +def _context_record( + module: Any, + torch: Any, + tokenizer: Any, + binding: Mapping[str, Any], + context: int, + prompt: Sequence[int], + witness: Mapping[str, Any], +) -> dict[str, Any]: + prefill = _prefill(module, torch, prompt) + require_prefill_token_match(context, prefill["bonus"], witness) + tokens: list[int] = [prefill["bonus"]] + rounds: list[dict[str, Any]] = [] + seed = prefill["bonus"] + position = len(prompt) + stopped: int | None = None + replay: dict[str, Any] | None = None + while stopped is None and len(tokens) < binding["token_budget"]: + raw = _read_round(module, _round_inputs(torch, seed, position)) + if replay is None: + replay = raw + decision = reconcile_k2_round( + raw["candidates"], + raw["target_greedy"], + raw["match_count"], + raw["bonus"], + raw["state_probe"], + position, + binding["token_budget"] - len(tokens), + binding["stop_tokens"], + ) + if not decision["valid"]: + raise OracleError(f"context {context} produced an invalid K=2 round") + tokens.extend(decision["committed"]) + rounds.append({"kv_witness": _kv_witness(torch, module), **raw, **decision}) + stopped = decision["stop_token"] if decision["stopped"] else None + seed = decision["next_seed"] + position = decision["next_position"] + return { + "accepted_prefix": [entry["accepted_drafts"] for entry in rounds], + "bonus_accounting": [entry["next_seed"] for entry in rounds], + "decoded_text": tokenizer.decode(tokens), + "kv_witnesses": [entry["kv_witness"] for entry in rounds], + "reset_replay": replay, + "rounds": rounds, + "selected_logits": [entry["selected"] for entry in rounds], + "stop_handling": { + "discarded": [entry["discarded"] for entry in rounds], + "stop_token": stopped, + }, + "target_prefill": witness, + "useful_tokens": tokens, + } + + +def run_portable_eager_authority( + binding: Mapping[str, Any] +) -> dict[str, dict[str, Any]]: + """Replay the D8 K=2 graph in eager mode; requires torch and the D8 module.""" + import torch + + from executorch.examples.models.gemma4.export_speculative import ( + build_k2_round_program, + ) + from transformers import AutoTokenizer + + checkpoints = binding["checkpoints"] + tokenizer = AutoTokenizer.from_pretrained(str(checkpoints["target"])) + program = build_k2_round_program( + checkpoints["target"], + checkpoints["assistant"], + max_seq_len=binding["max_context_length"], + max_input_len=K2_MAX_INPUT_LEN, + ) + records: dict[str, dict[str, Any]] = {} + for context in binding["contexts"]: + prompt = [(index % (K2_VOCAB_SIZE - 1)) + 1 for index in range(context)] + module = program.module() + record = _context_record( + module, + torch, + tokenizer, + binding, + context, + prompt, + binding["target_prefill"][str(context)], + ) + replayed = _read_round( + program.module(), _round_inputs(torch, record["useful_tokens"][0], context) + ) + if replayed != record["reset_replay"]: + raise OracleError(f"context {context} reset replay is not deterministic") + records[str(context)] = record + return records + + +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description=( + "Generate a Gemma 4 MTP/spec oracle bound to target-only eager " + "prefill evidence; it does not independently validate the shared model" + ) + ) + parser.add_argument("--oracle-binding-manifest", type=Path) + parser.add_argument("--production-mtp-manifest", type=Path) + parser.add_argument("--target-checkpoint", type=Path) + parser.add_argument("--assistant-checkpoint", type=Path) + parser.add_argument("--combined-runtime-root", type=Path, required=True) + parser.add_argument("--authority", required=True) + parser.add_argument("--contexts", required=True) + parser.add_argument("--output", type=Path, required=True) + return parser + + +def main(argv: Sequence[str] | None = None) -> int: + parser = _build_parser() + args = parser.parse_args(argv) + direct_inputs = ( + args.production_mtp_manifest, + args.target_checkpoint, + args.assistant_checkpoint, + ) + if args.oracle_binding_manifest is not None: + if any(value is not None for value in direct_inputs): + parser.error( + "--oracle-binding-manifest cannot be combined with direct production inputs" + ) + binding = build_oracle_binding( + args.oracle_binding_manifest, + None, + args.authority, + args.contexts, + args.output, + combined_runtime_root=args.combined_runtime_root, + ) + else: + if any(value is None for value in direct_inputs): + parser.error( + "direct production mode requires --production-mtp-manifest, " + "--target-checkpoint, and --assistant-checkpoint" + ) + binding = build_production_oracle_binding( + args.production_mtp_manifest, + args.target_checkpoint, + args.assistant_checkpoint, + args.authority, + args.contexts, + args.output, + combined_runtime_root=args.combined_runtime_root, + ) + document = assemble_oracle_document(binding, run_portable_eager_authority(binding)) + if binding.get("closure_state") == "full" and not production_oracle_is_acceptable( + document + ): + raise OracleError("generated oracle failed production validation") + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text( + json.dumps(document, indent=2, sort_keys=True) + "\n", encoding="utf-8" + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/models/gemma4/tests/targets.bzl b/examples/models/gemma4/tests/targets.bzl index fd1d12ecce4..32c64133fdf 100644 --- a/examples/models/gemma4/tests/targets.bzl +++ b/examples/models/gemma4/tests/targets.bzl @@ -5,6 +5,23 @@ def define_common_targets(is_fbcode = False): if not is_fbcode: return + # `mtp_export_lib` is D8-owned, so every target naming it needs D8 landed first. + fbcode_target(_kind = runtime.python_test, + name = "test_eagle_combined_round", + srcs = ["test_eagle_combined_round.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:op_registry", + "//executorch/examples/models/gemma4:mtp_export_lib", + "//executorch/exir:lib", + "//executorch/extension/pybindings:portable_lib", + ], + ) + fbcode_target(_kind = runtime.python_test, name = "test_speech_transform", srcs = ["test_speech_transform.py"], @@ -17,12 +34,18 @@ def define_common_targets(is_fbcode = False): fbcode_target(_kind = runtime.python_test, name = "test_export_partitioners", srcs = ["test_export_partitioners.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:op_registry", "//executorch/backends/vulkan/partitioner:vulkan_partitioner", + "//executorch/examples/models/gemma4:mtp_export_lib", "//executorch/examples/models/gemma4:webgpu_support", "//executorch/exir:lib", + "//executorch/extension/pybindings:portable_lib", ], ) @@ -92,3 +115,88 @@ def define_common_targets(is_fbcode = False): "//executorch/examples/models/gemma4:target_prefill_producer", ], ) + + fbcode_target(_kind = runtime.python_test, + name = "test_export_assistant_webgpu_artifacts", + srcs = ["test_export_assistant_webgpu_artifacts.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/examples/models/gemma4:mtp_export_lib", + "//executorch/examples/models/gemma4:webgpu_support", + "//executorch/exir:lib", + "//executorch/extension/pybindings:portable_lib", + ], + ) + + # The oracle generator ships in srcs so the suite carries its own reference. + fbcode_target(_kind = runtime.python_test, + name = "test_mtp_spec_oracle", + srcs = [ + "generate_mtp_spec_oracle.py", + "test_mtp_spec_oracle.py", + ], + 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 = [ + "//caffe2:torch", + "//executorch/examples/models/gemma4:mtp_export_lib", + "//executorch/examples/models/gemma4:target_prefill_producer", + "//executorch/examples/models/gemma4:webgpu_support", + "//executorch/exir:lib", + ], + typing = True, + ) + + fbcode_target(_kind = runtime.python_binary, + name = "generate_mtp_spec_oracle", + srcs = ["generate_mtp_spec_oracle.py"], + main_function = "executorch.examples.models.gemma4.tests.generate_mtp_spec_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 = [ + "//caffe2:torch", + "//executorch/examples/models/gemma4:mtp_export_lib", + "//executorch/examples/models/gemma4:target_prefill_producer", + "//executorch/examples/models/gemma4:webgpu_support", + "//executorch/exir:lib", + ], + typing = True, + ) + + fbcode_target(_kind = runtime.python_test, + name = "test_webgpu_spec_contract", + srcs = ["test_webgpu_spec_contract.py"], + deps = [ + "//executorch/examples/models/gemma4:webgpu_support", + ], + typing = True, + ) + + # The OSS closure gate scans the whole checkout; the native CI script runs it. + fbcode_target(_kind = runtime.python_library, + name = "test_oss_source_closure", + srcs = ["test_oss_source_closure.py"], + typing = True, + ) + + fbcode_target(_kind = runtime.cxx_test, + name = "test_gemma4_spec_runner_contract", + srcs = ["test_gemma4_spec_runner_contract.cpp"], + deps = [ + "//executorch/examples/models/gemma4:gemma4_spec_runner", + ], + ) diff --git a/examples/models/gemma4/tests/test_eagle_combined_round.py b/examples/models/gemma4/tests/test_eagle_combined_round.py new file mode 100644 index 00000000000..e2faf3613b0 --- /dev/null +++ b/examples/models/gemma4/tests/test_eagle_combined_round.py @@ -0,0 +1,1226 @@ +# 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. + +# pyre-strict + +from __future__ import annotations + +import contextlib +import dataclasses +import unittest +from collections.abc import Iterator +from pathlib import Path + +import executorch.extension.llm.custom_ops.custom_ops # noqa: F401 + +import torch + +from executorch.examples.models.gemma4.eagle_webgpu_round import ( + _expected_mutation_contract, + _rewrite_negative_select_as_symint, + export_k2_round_program, + Gemma4K2Target, + K2GPUResidentRound, + K2LongestPrefixSelector, + validate_k2_round_abi, +) +from executorch.examples.models.gemma4.export_speculative import ( + _lower_k2_round, + build_k2_round_program, +) +from executorch.examples.models.gemma4.webgpu_partitioner import ( + _is_official_qat_topk, + _is_official_qat_unique_scatter, + build_webgpu_partitioner, + mtp_extra_op_features, + rewrite_certified_unique_scatter, +) +from executorch.exir.dialects._ops import ops as exir_ops +from torch.export.graph_signature import ( + InputKind, + OutputKind, + OutputSpec, + TensorArgument, +) + + +_ABI_MAX_INPUT_LEN = 8 +_ABI_MAX_DONOR_LEN = 8960 +_ABI_MUTATION_COUNT = 31 +_ABI_SDPA_COUNT = 43 +_ABI_ARGMAX_COUNT = 3 +_ABI_TOPK_COUNT = 2 +_ABI_SCATTER_COUNT = 2 + +_FULL_DONOR_LAYER = 14 +_SLIDING_DONOR_LAYER = 13 +_DONOR_LAYER_COUNT = 16 + +_DONOR_LENGTHS = (2, 511, 512, 513, 514, 8960) + + +def _mutation_buffer_names( + count: int = _ABI_MUTATION_COUNT, + seed_names: tuple[str, ...] = ("seed_feature",), +) -> tuple[str, ...]: + filler = tuple(f"donor_cache_{index}" for index in range(count - len(seed_names))) + return seed_names + filler + + +class _K2AbiFixture(torch.nn.Module): + """Smallest module whose exported graph reproduces the `k2_round` ABI.""" + + def __init__( + self, + *, + max_input_len: int = _ABI_MAX_INPUT_LEN, + max_donor_len: int = _ABI_MAX_DONOR_LEN, + mutation_names: tuple[str, ...] | None = None, + sdpa_count: int = _ABI_SDPA_COUNT, + argmax_count: int = _ABI_ARGMAX_COUNT, + topk_count: int = _ABI_TOPK_COUNT, + scatter_count: int = _ABI_SCATTER_COUNT, + ) -> None: + super().__init__() + self.max_donor_len = max_donor_len + self.sdpa_count = sdpa_count + self.argmax_count = argmax_count + self.topk_count = topk_count + self.scatter_count = scatter_count + self.mutation_names: tuple[str, ...] = ( + _mutation_buffer_names() if mutation_names is None else mutation_names + ) + self.register_buffer( + "round_tail", torch.zeros((1, max_input_len - 3)), persistent=False + ) + self.register_buffer( + "donor_pool", torch.zeros((1, max_donor_len, 1, 1)), persistent=False + ) + self.register_buffer( + "attn_mask", torch.zeros((1, max_donor_len)), persistent=False + ) + self.register_buffer("query", torch.zeros((1, 1, 1, 1)), persistent=False) + self.register_buffer("logits", torch.zeros((1, 1, 8)), persistent=False) + self.register_buffer("scatter_row", torch.zeros((1, 4)), persistent=False) + self.register_buffer( + "scatter_index", torch.zeros((1, 2), dtype=torch.int64), persistent=False + ) + for name in self.mutation_names: + shape = (1, 1, 1, 1536) if name == "seed_feature" else (1, 1, 1, 1) + self.register_buffer(name, torch.zeros(shape)) + + def forward( + self, + input_ids: torch.Tensor, + input_pos: torch.Tensor, + is_round: torch.Tensor, + donor_length: torch.Tensor, + ) -> tuple[torch.Tensor, ...]: + donor_k = donor_length[0, 0].item() + torch._check_is_size(donor_k) + torch._check(donor_k >= 2) + torch._check(donor_k <= self.max_donor_len) + donor = self.donor_pool.narrow(1, 0, donor_k) + mask = self.attn_mask.narrow(1, 0, donor_k) + live = ( + input_ids.to(torch.float32).sum() + + input_pos.to(torch.float32).sum() + + is_round.to(torch.float32).sum() + ) + attention = self.query + live + for _ in range(self.sdpa_count): + attention = torch.ops.llama.custom_sdpa.default( + attention, donor, donor, 0, mask, 0.0, False, 1.0 + ) + scores = self.logits + attention.reshape(1, 1, 1) + for _ in range(self.topk_count): + scores = scores + torch.topk(scores, 2, dim=-1).values.sum() + scattered = self.scatter_row + for _ in range(self.scatter_count): + scattered = scattered.scatter( + -1, self.scatter_index, scores.reshape(1, -1)[:, :2] + ) + greedy: list[torch.Tensor] = [] + probe = scores + for _ in range(self.argmax_count): + greedy.append(torch.argmax(probe, dim=-1)) + probe = probe + 1.0 + next_feature = (live + attention.reshape(1, 1)).reshape(1, 1, 1).expand( + 1, 1, 1536 + ) + torch.ops.llama.update_cache.default( + next_feature.unsqueeze(2), self.seed_feature, 0 + ) + cache_value = live.reshape(1, 1, 1, 1) + for name in self.mutation_names[1:]: + torch.ops.llama.update_cache.default(cache_value, getattr(self, name), 0) + head = greedy[0] + bonus = head + 1 + for value in greedy[1:]: + bonus = bonus + value + seed = self.get_buffer("seed_feature") + return ( + torch.cat((head, head), dim=1), + torch.cat((head, head, head), dim=1), + head.reshape(1), + bonus, + (seed[..., 0] + scattered.sum() + donor.sum()).reshape(1, 1), + ) + + +def _export_abi_fixture(**overrides: int) -> torch.export.ExportedProgram: + fixture = _K2AbiFixture(**overrides).eval() + program = export_k2_round_program( + # pyre-ignore[6]: the ABI fixture duck-types `K2GPUResidentRound`. + fixture, + _ABI_MAX_INPUT_LEN, + ) + expected = _expected_mutation_contract(_ABI_MAX_DONOR_LEN) + signature = program.graph_signature + nodes = {node.name: node for node in program.graph.nodes} + placeholders = {target: name for name, target in signature.inputs_to_buffers.items()} + mutation_specs: list[OutputSpec] = [] + for source_target, record in zip(fixture.mutation_names, expected): + name = placeholders[source_target] + nodes[name].meta["val"] = torch.empty( + tuple(record["shape"]), dtype=torch.float32, device="meta" + ) + mutation_specs.append( + OutputSpec( + kind=OutputKind.BUFFER_MUTATION, + arg=TensorArgument(name=name), + target=record["logicalTarget"], + ) + ) + signature.output_specs[:] = mutation_specs + [ + spec for spec in signature.output_specs if spec.kind == OutputKind.USER_OUTPUT + ] + return program + + +def _range_bounds(program: torch.export.ExportedProgram) -> set[tuple[int, int]]: + bounds: set[tuple[int, int]] = set() + for value in program.range_constraints.values(): + try: + bounds.add((int(value.lower), int(value.upper))) + except (OverflowError, TypeError, ValueError): + continue + return bounds + + +@contextlib.contextmanager +def _reversible_abi_edit(program: torch.export.ExportedProgram) -> Iterator[None]: + """Undo signature / range / node-meta edits so one export can serve many mutants.""" + signature = program.graph_signature + saved_inputs = list(signature.input_specs) + saved_outputs = list(signature.output_specs) + saved_ranges = dict(program.range_constraints) + saved_vals = {node.name: node.meta.get("val") for node in program.graph.nodes} + try: + yield + finally: + signature.input_specs[:] = saved_inputs + signature.output_specs[:] = saved_outputs + program.range_constraints.clear() + program.range_constraints.update(saved_ranges) + for node in program.graph.nodes: + if node.name in saved_vals: + node.meta["val"] = saved_vals[node.name] + + +class K2RoundAbiTest(unittest.TestCase): + """`validate_k2_round_abi` pins the exact `k2_round` graph ABI.""" + + program: torch.export.ExportedProgram + + @classmethod + def setUpClass(cls) -> None: + cls.program = _export_abi_fixture() + + def _validate(self, program: torch.export.ExportedProgram) -> object: + return validate_k2_round_abi( + program, + max_input_len=_ABI_MAX_INPUT_LEN, + max_donor_len=_ABI_MAX_DONOR_LEN, + ) + + def _nodes(self) -> dict[str, torch.fx.Node]: + return {node.name: node for node in self.program.graph.nodes} + + def _mutation_indices(self) -> list[int]: + return [ + index + for index, spec in enumerate(self.program.graph_signature.output_specs) + if spec.kind == OutputKind.BUFFER_MUTATION + ] + + def test_abi_evidence_reports_the_documented_census(self) -> None: + evidence = self._validate(self.program) + self.assertEqual( + set(evidence), + { + "bufferMutationCount", + "donorViewOrder", + "inputOrder", + "mutationOrder", + "operatorCounts", + "outputOrder", + "seedMutationCount", + "stateAlias", + }, + ) + self.assertEqual(evidence["bufferMutationCount"], 31) + self.assertEqual( + evidence["operatorCounts"], + { + "aten.argmax.default": 3, + "aten.scatter.src": 2, + "aten.topk.default": 2, + "llama.custom_sdpa.default": 43, + "llama.update_cache.default": 31, + }, + ) + self.assertEqual(evidence["seedMutationCount"], 1) + + def test_user_inputs_are_ordered_ids_pos_round_donor(self) -> None: + self.assertEqual( + tuple(self.program.graph_signature.user_inputs), + ("input_ids", "input_pos", "is_round", "donor_length"), + ) + + def test_user_input_shapes_and_dtypes_are_exact(self) -> None: + nodes = self._nodes() + for name in ("input_ids", "input_pos", "is_round", "donor_length"): + self.assertEqual(nodes[name].op, "placeholder", name) + self.assertEqual(nodes[name].meta["val"].dtype, torch.int64, name) + self.assertEqual(len(nodes["input_ids"].meta["val"].shape), 2) + self.assertEqual(nodes["input_ids"].meta["val"].shape[0], 1) + self.assertEqual(len(nodes["input_pos"].meta["val"].shape), 1) + self.assertEqual(tuple(nodes["is_round"].meta["val"].shape), (1,)) + self.assertEqual(tuple(nodes["donor_length"].meta["val"].shape), (1, 1)) + + def test_input_order_permutation_is_rejected(self) -> None: + specs = self.program.graph_signature.input_specs + first, second = ( + index + for index, spec in enumerate(specs) + if spec.kind == InputKind.USER_INPUT + and spec.arg.name in ("input_ids", "input_pos") + ) + with _reversible_abi_edit(self.program): + specs[first], specs[second] = specs[second], specs[first] + with self.assertRaisesRegex(ValueError, "user-input order mismatch"): + self._validate(self.program) + + def test_input_dtype_and_shape_regressions_are_rejected(self) -> None: + mutants = { + "is_round": torch.zeros((1,), dtype=torch.int32), + "donor_length": torch.zeros((1,), dtype=torch.int64), + "input_ids": torch.zeros((2, 3), dtype=torch.int64), + } + nodes = self._nodes() + for name, mutant in mutants.items(): + with self.subTest(name=name), _reversible_abi_edit(self.program): + nodes[name].meta["val"] = mutant + with self.assertRaisesRegex(ValueError, f"K=2 {name}"): + self._validate(self.program) + + def test_sequence_dimension_is_one_shared_symbol(self) -> None: + nodes = self._nodes() + input_ids = nodes["input_ids"].meta["val"] + input_pos = nodes["input_pos"].meta["val"] + self.assertEqual(str(input_ids.shape[1]), str(input_pos.shape[0])) + self.assertNotEqual(str(input_ids.shape[1]), str(input_ids.shape[0])) + with _reversible_abi_edit(self.program): + nodes["input_pos"].meta["val"] = torch.zeros((3,), dtype=torch.int64) + with self.assertRaisesRegex(ValueError, "dynamic dimensions differ"): + self._validate(self.program) + + def test_user_outputs_are_ordered_and_typed(self) -> None: + nodes = self._nodes() + outputs = list(self.program.graph_signature.user_outputs) + self.assertEqual(len(outputs), 5) + expected = ( + ((1, 2), torch.int64), + ((1, 3), torch.int64), + ((1,), torch.int64), + ((1, 1), torch.int64), + ((1, 1), torch.float32), + ) + for name, (shape, dtype) in zip(outputs, expected): + value = nodes[str(name)].meta["val"] + self.assertEqual(tuple(value.shape), shape, name) + self.assertEqual(value.dtype, dtype, name) + + def test_output_order_permutation_is_rejected(self) -> None: + specs = self.program.graph_signature.output_specs + user = [ + index + for index, spec in enumerate(specs) + if spec.kind == OutputKind.USER_OUTPUT + ] + with _reversible_abi_edit(self.program): + first, last = user[0], user[-1] + specs[first], specs[last] = specs[last], specs[first] + with self.assertRaisesRegex(ValueError, "K=2 candidates"): + self._validate(self.program) + + def test_state_probe_must_stay_float32(self) -> None: + probe = str(list(self.program.graph_signature.user_outputs)[-1]) + nodes = self._nodes() + with _reversible_abi_edit(self.program): + nodes[probe].meta["val"] = torch.zeros((1, 1), dtype=torch.int64) + with self.assertRaisesRegex(ValueError, "K=2 state_probe dtype"): + self._validate(self.program) + + def test_missing_user_output_is_rejected(self) -> None: + specs = self.program.graph_signature.output_specs + with _reversible_abi_edit(self.program): + dropped = next( + index + for index, spec in enumerate(specs) + if spec.kind == OutputKind.USER_OUTPUT + ) + del specs[dropped] + with self.assertRaisesRegex(ValueError, "user-output count mismatch"): + self._validate(self.program) + + def test_mutation_census_is_thirty_one_with_one_seed(self) -> None: + specs = self.program.graph_signature.output_specs + mutations = self._mutation_indices() + targets = [str(specs[index].target) for index in mutations] + self.assertEqual(len(mutations), 31) + self.assertEqual( + [target for target in targets if target.endswith("seed_feature")], + ["seed_feature"], + ) + self.assertEqual(len(set(targets)), 31) + + def test_dropped_mutation_is_rejected(self) -> None: + specs = self.program.graph_signature.output_specs + with _reversible_abi_edit(self.program): + del specs[self._mutation_indices()[-1]] + with self.assertRaisesRegex(ValueError, "output-spec order mismatch"): + self._validate(self.program) + + def test_second_seed_feature_mutation_is_rejected(self) -> None: + specs = self.program.graph_signature.output_specs + with _reversible_abi_edit(self.program): + victim = next( + index + for index in self._mutation_indices() + if not str(specs[index].target).endswith("seed_feature") + ) + specs[victim] = dataclasses.replace( + specs[victim], target="assistant.seed_feature" + ) + with self.assertRaisesRegex(ValueError, "mutation target order mismatch"): + self._validate(self.program) + + def test_duplicate_mutation_target_is_rejected(self) -> None: + specs = self.program.graph_signature.output_specs + with _reversible_abi_edit(self.program): + non_seed = [ + index + for index in self._mutation_indices() + if not str(specs[index].target).endswith("seed_feature") + ] + specs[non_seed[1]] = dataclasses.replace( + specs[non_seed[1]], target=str(specs[non_seed[0]].target) + ) + with self.assertRaisesRegex(ValueError, "mutation target order mismatch"): + self._validate(self.program) + + def test_operator_census_is_exact(self) -> None: + mutants = { + "llama.custom_sdpa.default": {"sdpa_count": _ABI_SDPA_COUNT - 1}, + "aten.argmax.default": {"argmax_count": _ABI_ARGMAX_COUNT - 1}, + "aten.topk.default": {"topk_count": _ABI_TOPK_COUNT - 1}, + "aten.scatter.src": {"scatter_count": _ABI_SCATTER_COUNT - 1}, + } + for target, override in mutants.items(): + with self.subTest(target=target): + program = _export_abi_fixture(**override) + with self.assertRaisesRegex(ValueError, f"{target} count mismatch"): + self._validate(program) + + def test_both_dynamic_ranges_must_be_present(self) -> None: + bounds = _range_bounds(self.program) + self.assertIn((1, _ABI_MAX_INPUT_LEN), bounds) + self.assertIn((2, _ABI_MAX_DONOR_LEN), bounds) + for missing in ((1, _ABI_MAX_INPUT_LEN), (2, _ABI_MAX_DONOR_LEN)): + with self.subTest(missing=missing), _reversible_abi_edit(self.program): + constraints = self.program.range_constraints + for symbol, value in list(constraints.items()): + try: + current = (int(value.lower), int(value.upper)) + except (OverflowError, TypeError, ValueError): + continue + if current == missing: + del constraints[symbol] + with self.assertRaisesRegex(ValueError, "missing dynamic range"): + self._validate(self.program) + + +class _DonorKVCache(torch.nn.Module): + def __init__(self, offset: float, max_seq_len: int) -> None: + super().__init__() + base = torch.arange(max_seq_len * 8, dtype=torch.float32).reshape( + 1, max_seq_len, 2, 4 + ) + self.register_buffer("k_cache", base + offset) + self.register_buffer("v_cache", base + offset + 0.5) + + +class _DonorAttention(torch.nn.Module): + def __init__( + self, + *, + is_donor: bool = False, + is_sliding: bool = False, + kv_cache: _DonorKVCache | None = None, + ) -> None: + super().__init__() + self.is_kv_donor_layer = is_donor + self.is_sliding = is_sliding + self.kv_cache = kv_cache + + +class _DonorLayer(torch.nn.Module): + def __init__(self, self_attn: _DonorAttention) -> None: + super().__init__() + self.self_attn = self_attn + + +class _DonorTextModel(torch.nn.Module): + def __init__( + self, + *, + full_index: int = _FULL_DONOR_LAYER, + sliding_index: int = _SLIDING_DONOR_LAYER, + max_seq_len: int = 32, + drop_full_cache: bool = False, + ) -> None: + super().__init__() + layers: list[_DonorLayer] = [] + for index in range(_DONOR_LAYER_COUNT): + if index == full_index: + layers.append( + _DonorLayer( + _DonorAttention( + is_donor=True, + is_sliding=False, + kv_cache=( + None + if drop_full_cache + else _DonorKVCache(0.0, max_seq_len) + ), + ) + ) + ) + elif index == sliding_index: + layers.append( + _DonorLayer( + _DonorAttention( + is_donor=True, + is_sliding=True, + kv_cache=_DonorKVCache(1000.0, max_seq_len), + ) + ) + ) + else: + layers.append(_DonorLayer(_DonorAttention())) + self.self_decoder = torch.nn.Module() + self.self_decoder.layers = torch.nn.ModuleList(layers) + + +class DonorTopologyTest(unittest.TestCase): + """`Gemma4K2Target` binds exactly one full donor (14) and one sliding donor (13).""" + + def test_official_topology_binds_layers_thirteen_and_fourteen(self) -> None: + target = Gemma4K2Target(_DonorTextModel()) + self.assertEqual(target.full_donor_index, 14) + self.assertEqual(target.sliding_donor_index, 13) + + def test_shifted_donor_layers_are_rejected(self) -> None: + with self.assertRaisesRegex(ValueError, "donor layer mismatch"): + Gemma4K2Target(_DonorTextModel(full_index=15)) + with self.assertRaisesRegex(ValueError, "donor layer mismatch"): + Gemma4K2Target(_DonorTextModel(sliding_index=12)) + + def test_duplicate_or_absent_donors_are_rejected(self) -> None: + with self.assertRaisesRegex(ValueError, "donor topology mismatch"): + Gemma4K2Target(_DonorTextModel(sliding_index=_FULL_DONOR_LAYER)) + empty = _DonorTextModel() + for layer in empty.self_decoder.layers: + layer.self_attn.is_kv_donor_layer = False + with self.assertRaisesRegex(ValueError, "donor topology mismatch"): + Gemma4K2Target(empty) + + def test_donor_without_kv_cache_is_rejected(self) -> None: + with self.assertRaisesRegex(ValueError, "donor has no KV cache"): + Gemma4K2Target(_DonorTextModel(drop_full_cache=True)) + + +class DonorViewTest(unittest.TestCase): + """Donor views are `cache[:, :length]` in BHKD order fk / fv / sk / sv.""" + + def test_views_are_prefix_slices_in_bhkd_layout(self) -> None: + text_model = _DonorTextModel(max_seq_len=32) + target = Gemma4K2Target(text_model) + layers = text_model.self_decoder.layers + full = layers[_FULL_DONOR_LAYER].self_attn.kv_cache + sliding = layers[_SLIDING_DONOR_LAYER].self_attn.kv_cache + for length in (2, 17, 32): + with self.subTest(length=length): + views = target.donor_views(torch.tensor([[length]], dtype=torch.int64)) + expected = ( + full.k_cache[:, :length].permute(0, 2, 1, 3), + full.v_cache[:, :length].permute(0, 2, 1, 3), + sliding.k_cache[:, :length].permute(0, 2, 1, 3), + sliding.v_cache[:, :length].permute(0, 2, 1, 3), + ) + self.assertEqual(len(views), 4) + for index, (view, reference) in enumerate(zip(views, expected)): + self.assertEqual(tuple(view.shape), (1, 2, length, 4), index) + self.assertTrue(torch.equal(view, reference), index) + + def test_view_order_is_not_interchangeable(self) -> None: + target = Gemma4K2Target(_DonorTextModel(max_seq_len=32)) + full_k, full_v, sliding_k, sliding_v = target.donor_views( + torch.tensor([[8]], dtype=torch.int64) + ) + for left, right in ( + (full_k, full_v), + (full_k, sliding_k), + (sliding_k, sliding_v), + ): + self.assertFalse(torch.equal(left, right)) + + def test_donor_length_below_two_is_rejected(self) -> None: + target = Gemma4K2Target(_DonorTextModel(max_seq_len=32)) + with self.assertRaises(RuntimeError): + target.donor_views(torch.tensor([[1]], dtype=torch.int64)) + + +class _RecordingTarget(torch.nn.Module): + """Stand-in target that reports the donor length it was handed.""" + + def __init__(self, *, donor_shrink: int = 0, head_dim: int = 2) -> None: + super().__init__() + self.donor_shrink = donor_shrink + self.head_dim = head_dim + self.donor_lengths: list[int] = [] + + def donor_views( + self, donor_length: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + length = int(donor_length[0, 0].item()) + self.donor_lengths.append(length) + width = length - self.donor_shrink + return ( + torch.full((1, 1, width, self.head_dim), 0.0), + torch.full((1, 1, width, self.head_dim), 1.0), + torch.full((1, 1, width, self.head_dim), 2.0), + torch.full((1, 1, width, self.head_dim), 3.0), + ) + + def forward( + self, + input_ids: torch.Tensor, + input_pos: torch.Tensor, + is_round: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + del input_ids, input_pos, is_round + return ( + torch.tensor([[21, 22, 23]], dtype=torch.long), + torch.tensor([[[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]]]), + ) + + +class _ZeroEmbedding(torch.nn.Module): + def forward(self, tokens: torch.Tensor) -> torch.Tensor: + return torch.zeros((*tokens.shape, 2), dtype=torch.float32) + + +class _TokenEmbedding(torch.nn.Module): + def forward(self, tokens: torch.Tensor) -> torch.Tensor: + return tokens.to(torch.float32).unsqueeze(-1).expand(*tokens.shape, 2) + + +class _ScaleDroppingEmbedding(_TokenEmbedding): + """Mutant: the caller-applied scale is cancelled on the second draft step.""" + + def __init__(self, embed_scale: float) -> None: + super().__init__() + self.embed_scale = embed_scale + self.calls = 0 + + def forward(self, tokens: torch.Tensor) -> torch.Tensor: + embedding = super().forward(tokens) + self.calls += 1 + return embedding if self.calls == 1 else embedding / self.embed_scale + + +class _RecordingAssistant(torch.nn.Module): + def __init__(self, first_token: int = 21) -> None: + super().__init__() + self.first_token = first_token + self.inputs: list[torch.Tensor] = [] + self.positions: list[list[list[int]]] = [] + self.donor_shapes: list[tuple[tuple[int, ...], ...]] = [] + self.donor_markers: list[tuple[float, ...]] = [] + + def forward( + self, + inputs: torch.Tensor, + position_ids: torch.Tensor, + full_k: torch.Tensor, + full_v: torch.Tensor, + sliding_k: torch.Tensor, + sliding_v: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + donors = (full_k, full_v, sliding_k, sliding_v) + self.inputs.append(inputs.clone()) + self.positions.append(position_ids.tolist()) + self.donor_shapes.append(tuple(tuple(donor.shape) for donor in donors)) + self.donor_markers.append( + tuple(float(donor.reshape(-1)[0].item()) for donor in donors) + ) + logits = torch.zeros((1, 1, 32), dtype=torch.float32) + logits[..., self.first_token + len(self.inputs) - 1] = 1.0 + return logits, inputs[..., :2] + + +class _EagerK2Round(K2GPUResidentRound): + """Keeps eager rounds off `llama.update_cache` so its dep is exercised once.""" + + def update_seed_feature(self, next_feature: torch.Tensor) -> torch.Tensor: + return next_feature.unsqueeze(2) + + +def _build_round( + target: torch.nn.Module, + embedding: torch.nn.Module, + assistant: torch.nn.Module, + *, + embed_scale: float = 1.0, + round_class: type[K2GPUResidentRound] = _EagerK2Round, +) -> K2GPUResidentRound: + return round_class( + target, + embedding, + assistant, + hidden_size=2, + max_input_len=3, + max_donor_len=_ABI_MAX_DONOR_LEN, + embed_scale=embed_scale, + ) + + +def _run_round( + module: K2GPUResidentRound, donor_length: int, *, is_round: int = 1 +) -> tuple[torch.Tensor, ...]: + return module( + torch.tensor([[10, 0, 0]], dtype=torch.long), + torch.arange(donor_length, donor_length + 3, dtype=torch.long), + torch.tensor([is_round], dtype=torch.long), + torch.tensor([[donor_length]], dtype=torch.long), + ) + + +class K2RoundStepTest(unittest.TestCase): + """K=2 assistant steps advance from `donor_length - 1` to `donor_length`.""" + + def _assert_advancing_positions_and_full_donors( + self, assistant: _RecordingAssistant, donor_length: int, head_dim: int = 2 + ) -> None: + self.assertEqual(len(assistant.positions), 2) + self.assertEqual( + assistant.positions, [[[donor_length - 1]], [[donor_length]]] + ) + expected_shapes = tuple((1, 1, donor_length, head_dim) for _ in range(4)) + self.assertEqual(assistant.donor_shapes, [expected_shapes, expected_shapes]) + self.assertEqual( + assistant.donor_markers, [(0.0, 1.0, 2.0, 3.0), (0.0, 1.0, 2.0, 3.0)] + ) + + def test_positions_advance_and_donors_keep_full_length(self) -> None: + for donor_length in _DONOR_LENGTHS: + with self.subTest(donor_length=donor_length): + target = _RecordingTarget() + assistant = _RecordingAssistant() + module = _build_round(target, _ZeroEmbedding(), assistant) + _run_round(module, donor_length) + self.assertEqual(target.donor_lengths, [donor_length]) + self._assert_advancing_positions_and_full_donors(assistant, donor_length) + + def test_donor_shrink_mutant_is_rejected(self) -> None: + for donor_length in (2, 512, 8960): + with self.subTest(donor_length=donor_length): + assistant = _RecordingAssistant() + module = _build_round( + _RecordingTarget(donor_shrink=1), _ZeroEmbedding(), assistant + ) + _run_round(module, donor_length) + with self.assertRaises(AssertionError): + self._assert_advancing_positions_and_full_donors( + assistant, donor_length + ) + + def test_position_oracle_rejects_a_shifted_donor_length(self) -> None: + donor_length = 512 + assistant = _RecordingAssistant() + module = _build_round(_RecordingTarget(), _ZeroEmbedding(), assistant) + _run_round(module, donor_length) + for shifted in (donor_length - 1, donor_length + 1): + with self.subTest(shifted=shifted): + with self.assertRaises(AssertionError): + self._assert_advancing_positions_and_full_donors( + assistant, shifted + ) + + +class K2EmbeddingScaleTest(unittest.TestCase): + """`embed_scale` multiplies both draft embeddings and must be finite / positive.""" + + def _assert_both_drafts_scaled( + self, assistant: _RecordingAssistant, embed_scale: float + ) -> None: + self.assertEqual(len(assistant.inputs), 2) + self.assertTrue( + torch.equal( + assistant.inputs[0][..., :2], torch.full((1, 1, 2), 10.0 * embed_scale) + ) + ) + self.assertTrue( + torch.equal( + assistant.inputs[1][..., :2], torch.full((1, 1, 2), 21.0 * embed_scale) + ) + ) + + def test_scale_is_applied_to_both_draft_embeddings(self) -> None: + for embed_scale in (0.5, 2.0): + with self.subTest(embed_scale=embed_scale): + assistant = _RecordingAssistant() + module = _build_round( + _RecordingTarget(), + _TokenEmbedding(), + assistant, + embed_scale=embed_scale, + ) + _run_round(module, 2) + self._assert_both_drafts_scaled(assistant, embed_scale) + + def test_dropping_the_scale_on_the_second_draft_is_rejected(self) -> None: + embed_scale = 0.5 + assistant = _RecordingAssistant() + module = _build_round( + _RecordingTarget(), + _ScaleDroppingEmbedding(embed_scale), + assistant, + embed_scale=embed_scale, + ) + _run_round(module, 2) + self.assertTrue( + torch.equal( + assistant.inputs[0][..., :2], torch.full((1, 1, 2), 10.0 * embed_scale) + ) + ) + with self.assertRaises(AssertionError): + self._assert_both_drafts_scaled(assistant, embed_scale) + + def test_non_finite_or_non_positive_scales_are_rejected(self) -> None: + for embed_scale in (float("nan"), float("inf"), float("-inf"), 0.0, -1.0): + with self.subTest(embed_scale=embed_scale): + with self.assertRaisesRegex(ValueError, "target embedding scale"): + _build_round( + _RecordingTarget(), + _TokenEmbedding(), + _RecordingAssistant(), + embed_scale=embed_scale, + ) + + def test_degenerate_round_dimensions_are_rejected(self) -> None: + for hidden_size, max_input_len, max_donor_len in ( + (0, 3, 2), + (2, 2, 2), + (2, 3, 1), + ): + with self.subTest(hidden_size=hidden_size, max_input_len=max_input_len): + with self.assertRaisesRegex(ValueError, "combined-round dimensions"): + K2GPUResidentRound( + _RecordingTarget(), + _TokenEmbedding(), + _RecordingAssistant(), + hidden_size=hidden_size, + max_input_len=max_input_len, + max_donor_len=max_donor_len, + embed_scale=1.0, + ) + + +class K2RoundOutputTest(unittest.TestCase): + """Round vs prefill selection of matches, bonus, feature and the state probe.""" + + def test_round_mode_emits_longest_prefix_evidence(self) -> None: + assistant = _RecordingAssistant() + module = _build_round(_RecordingTarget(), _TokenEmbedding(), assistant) + candidates, greedy, matches, bonus, probe = _run_round(module, 2) + self.assertEqual(candidates.tolist(), [[21, 22]]) + self.assertEqual(candidates.dtype, torch.int64) + self.assertEqual(greedy.tolist(), [[21, 22, 23]]) + self.assertEqual(matches.tolist(), [2]) + self.assertEqual(bonus.tolist(), [[23]]) + self.assertEqual(probe.dtype, torch.float32) + self.assertEqual(tuple(probe.shape), (1, 1)) + self.assertEqual(probe.tolist(), [[5.0]]) + + def test_rejected_drafts_fall_back_to_the_first_target_row(self) -> None: + assistant = _RecordingAssistant(first_token=5) + module = _build_round(_RecordingTarget(), _TokenEmbedding(), assistant) + candidates, _greedy, matches, bonus, probe = _run_round(module, 2) + self.assertEqual(candidates.tolist(), [[5, 6]]) + self.assertEqual(matches.tolist(), [0]) + self.assertEqual(bonus.tolist(), [[21]]) + self.assertEqual(probe.tolist(), [[1.0]]) + + def test_prefill_mode_zeroes_matches_and_takes_the_last_greedy(self) -> None: + assistant = _RecordingAssistant(first_token=5) + module = _build_round(_RecordingTarget(), _TokenEmbedding(), assistant) + _candidates, _greedy, matches, bonus, probe = _run_round(module, 2, is_round=0) + self.assertEqual(matches.tolist(), [0]) + self.assertEqual(bonus.tolist(), [[23]]) + self.assertEqual(probe.tolist(), [[5.0]]) + + def test_state_probe_reads_the_updated_seed_feature(self) -> None: + module = _build_round( + _RecordingTarget(), + _TokenEmbedding(), + _RecordingAssistant(), + round_class=K2GPUResidentRound, + ) + self.assertEqual(module.seed_feature.tolist(), [[[[0.0, 0.0]]]]) + probe = _run_round(module, 2)[-1] + self.assertEqual(module.seed_feature.tolist(), [[[[5.0, 6.0]]]]) + self.assertEqual(probe.tolist(), [[5.0]]) + + +class LongestPrefixSelectorTest(unittest.TestCase): + def test_prefix_length_drives_bonus_and_feature(self) -> None: + selector = K2LongestPrefixSelector() + features = torch.tensor([[[10.0, 11.0], [20.0, 21.0], [30.0, 31.0]]]) + greedy = torch.tensor([[90, 91, 92]]) + + for drafts, expected_count, expected_bonus, expected_feature in ( + ([80, 81], 0, 90, [10.0, 11.0]), + ([90, 81], 1, 91, [20.0, 21.0]), + ([90, 91], 2, 92, [30.0, 31.0]), + ([80, 91], 0, 90, [10.0, 11.0]), + ): + with self.subTest(drafts=drafts): + count, bonus, candidates, feature = selector( + torch.tensor([[7, *drafts]]), greedy, features + ) + self.assertEqual(count.tolist(), [expected_count]) + self.assertEqual(count.dtype, torch.int64) + self.assertEqual(bonus.tolist(), [expected_bonus]) + self.assertEqual(candidates.tolist(), [drafts]) + self.assertEqual(feature.tolist(), [expected_feature]) + + +class _NegativeSelectChains(torch.nn.Module): + def __init__(self, chains: int) -> None: + super().__init__() + self.chains = chains + + def forward(self, first: torch.Tensor, second: torch.Tensor) -> torch.Tensor: + total = torch.zeros((1,), dtype=torch.int64) + for source in (first, second)[: self.chains]: + total = total + source[-1].item() + return total + + +def _export_negative_select_chains(chains: int) -> torch.export.ExportedProgram: + positions = torch.arange(3, dtype=torch.int64) + return torch.export.export( + _NegativeSelectChains(chains), (positions, positions.clone()), strict=False + ) + + +class SelectAsSymintRewriteTest(unittest.TestCase): + """`compose_k2_round_program` demands exactly two negative-select rewrites.""" + + def test_rewrite_count_tracks_negative_select_chains(self) -> None: + for chains in (0, 1, 2): + with self.subTest(chains=chains): + self.assertEqual( + _rewrite_negative_select_as_symint( + _export_negative_select_chains(chains) + ), + chains, + ) + + +class _TinyRoundBoundFixture(torch.nn.Module): + def __init__(self, tail_width: int) -> None: + super().__init__() + self.register_buffer("round_tail", torch.zeros((1, tail_width))) + + +class ExportGuardTest(unittest.TestCase): + def test_round_tail_must_cover_the_declared_input_bound(self) -> None: + for max_input_len, tail_width in ((2, 0), (8, 4)): + with self.subTest(max_input_len=max_input_len): + with self.assertRaisesRegex(ValueError, "round input bound"): + export_k2_round_program( + # pyre-ignore[6]: the bound guard runs before any module use. + _TinyRoundBoundFixture(tail_width), + max_input_len, + ) + + +class Emb4TargetContractTest(unittest.TestCase): + """The K=2 target is pinned to `8da4w+emb4`, group size 128 and a 4-bit head.""" + + def _build(self, **overrides: object) -> None: + arguments: dict[str, object] = { + "max_seq_len": 8960, + "max_input_len": 512, + "text_quantize": "8da4w+emb4", + "assistant_quantize": "8da4w", + "assistant_lm_head_bits": 4, + "group_size": 128, + } + arguments.update(overrides) + build_k2_round_program( + Path("/nonexistent/gemma4-target"), + Path("/nonexistent/gemma4-assistant"), + # pyre-ignore[6]: parametrised guard arguments. + **arguments, + ) + + def test_emb8_and_other_target_quantizations_fail_closed(self) -> None: + for text_quantize in ("8da4w+emb8", "8da4w", "emb4", "8da4w+emb4 ", ""): + with self.subTest(text_quantize=text_quantize): + with self.assertRaisesRegex(ValueError, r"8da4w\+emb4"): + self._build(text_quantize=text_quantize) + + def test_group_size_is_pinned_to_128(self) -> None: + for group_size in (32, 64, 256): + with self.subTest(group_size=group_size): + with self.assertRaisesRegex(ValueError, "group size 128"): + self._build(group_size=group_size) + + def test_assistant_head_is_pinned_to_four_bits_and_8da4w(self) -> None: + with self.assertRaisesRegex(ValueError, "4-bit LM head"): + self._build(assistant_lm_head_bits=8) + with self.assertRaisesRegex(ValueError, "4-bit LM head"): + self._build(assistant_quantize="8da8w") + + def test_sequence_bounds_are_guarded(self) -> None: + with self.assertRaisesRegex(ValueError, "max_seq_len >= 514"): + self._build(max_seq_len=513) + with self.assertRaisesRegex(ValueError, "max_input_len"): + self._build(max_input_len=2) + with self.assertRaisesRegex(ValueError, "max_input_len"): + self._build(max_seq_len=1024, max_input_len=1025) + + def test_partitioner_rejects_non_emb4_targets(self) -> None: + with self.assertRaisesRegex(ValueError, "emb4"): + build_webgpu_partitioner(text_quantize="8da4w+emb8", mode="mtp") + + def test_partitioner_refuses_conflicting_compile_options(self) -> None: + with self.assertRaisesRegex(ValueError, "cannot override"): + build_webgpu_partitioner( + "8da4w+emb4", + mode="mtp", + compile_options={"require_dynamic_shapes": False}, + ) + with self.assertRaisesRegex(ValueError, "cannot override"): + build_webgpu_partitioner( + "8da4w+emb4", + mode="mtp", + compile_options={"skip_bool_tensors": True}, + ) + + def test_lowering_forwards_the_target_quantization_contract(self) -> None: + with self.assertRaisesRegex(ValueError, "emb4"): + _lower_k2_round( + _export_negative_select_chains(0), + external_constants_max_data_bytes=1024, + text_quantize="8da4w+emb8", + ) + + +class _TwoResidualChains(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.register_buffer( + "ordering", + torch.arange(262144, dtype=torch.float32).reshape(2048, 128), + persistent=False, + ) + self.register_buffer("output", torch.zeros((1, 1, 262144), dtype=torch.float32)) + + def _chain(self, scores: torch.Tensor) -> torch.Tensor: + _, indices = torch.topk(scores, 32, dim=-1) + destinations = ( + torch.nn.functional.embedding(indices, self.ordering) + .to(torch.long) + .view(1, 1, 4096) + ) + source = torch.ones((1, 1, 4096), dtype=torch.float32) + return self.output.scatter(-1, destinations, source) + + def forward( + self, first: torch.Tensor, second: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + return self._chain(first), self._chain(second) + + +class _DisconnectedResidualChains(_TwoResidualChains): + def _chain(self, scores: torch.Tensor) -> torch.Tensor: + values, _ = torch.topk(scores, 32, dim=-1) + destinations = torch.arange(4096).view(1, 1, 4096) + source = values.repeat_interleave(128, dim=-1) + return self.output.scatter(-1, destinations, source) + + +class _TopKOnly(torch.nn.Module): + def forward(self, scores: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + return torch.topk(scores, 32, dim=-1, largest=True, sorted=True) + + +class _ScatterOnly(torch.nn.Module): + def forward( + self, output: torch.Tensor, index: torch.Tensor, source: torch.Tensor + ) -> torch.Tensor: + return output.scatter(-1, index, source) + + +def _find_call( + program: torch.export.ExportedProgram, target: torch._ops.OpOverload +) -> torch.fx.Node: + return next( + node + for node in program.graph.nodes + if node.op == "call_function" and node.target == target + ) + + +def _export_chains(module: torch.nn.Module) -> torch.export.ExportedProgram: + scores = torch.zeros((1, 1, 2048), dtype=torch.float32) + return torch.export.export( + module, (scores, scores.clone()), strict=False + ).run_decompositions({}) + + +class MtpScatterRewriteTest(unittest.TestCase): + @unittest.skip("requires the final assistant residual topology") + def test_rewrite_is_scoped_to_two_certified_chains(self) -> None: + program = _export_chains(_TwoResidualChains()) + ordering = torch.arange(262144, dtype=torch.int64) + self.assertEqual(rewrite_certified_unique_scatter(program, ordering), 2) + targets = [ + node.target for node in program.graph.nodes if node.op == "call_function" + ] + self.assertEqual(targets.count(torch.ops.et_vk.scatter_src_unique.default), 2) + self.assertNotIn(torch.ops.aten.scatter.src, targets) + + def test_non_permutation_ordering_is_rejected(self) -> None: + duplicate = torch.arange(262144, dtype=torch.int64) + duplicate[-1] = duplicate[0] + with self.assertRaisesRegex(ValueError, "permutation"): + rewrite_certified_unique_scatter( + _export_chains(_TwoResidualChains()), duplicate + ) + + @unittest.skip("requires the final assistant residual topology") + def test_ordering_provenance_must_match_the_baked_buffer(self) -> None: + ordering = torch.arange(262144, dtype=torch.int64) + with self.assertRaisesRegex(ValueError, "token-ordering conversion mismatch"): + rewrite_certified_unique_scatter( + _export_chains(_TwoResidualChains()), ordering.flip(0) + ) + + def test_scatter_without_a_topk_ancestor_is_rejected(self) -> None: + ordering = torch.arange(262144, dtype=torch.int64) + with self.assertRaisesRegex(ValueError, "token-ordering conversion mismatch"): + rewrite_certified_unique_scatter( + _export_chains(_DisconnectedResidualChains()), ordering + ) + + def test_chain_count_other_than_two_is_rejected(self) -> None: + ordering = torch.arange(262144, dtype=torch.int64) + with self.assertRaisesRegex(ValueError, "residual topology mismatch"): + rewrite_certified_unique_scatter( + _export_chains(_TwoResidualChains()), ordering, expected_chains=3 + ) + with self.assertRaisesRegex(ValueError, "residual topology mismatch"): + rewrite_certified_unique_scatter( + _export_chains(_TwoResidualChains()), ordering, expected_chains=0 + ) + + +class MtpOpFeatureTest(unittest.TestCase): + def test_features_are_instance_scoped_and_exact(self) -> None: + from executorch.backends.vulkan.op_registry import vulkan_supported_ops + + before = dict(vulkan_supported_ops) + features = mtp_extra_op_features() + self.assertIn(exir_ops.edge.aten.topk.default, features) + self.assertIn(exir_ops.edge.et_vk.scatter_src_unique.default, features) + self.assertEqual(before, vulkan_supported_ops) + + def test_residual_routes_require_exact_full_shapes(self) -> None: + exact_topk = torch.export.export( + _TopKOnly(), + (torch.zeros((1, 1, 2048), dtype=torch.float32),), + strict=False, + ) + self.assertTrue( + _is_official_qat_topk(_find_call(exact_topk, torch.ops.aten.topk.default)) + ) + wrong_rank_topk = torch.export.export( + _TopKOnly(), + (torch.zeros((1, 2048), dtype=torch.float32),), + strict=False, + ) + self.assertFalse( + _is_official_qat_topk( + _find_call(wrong_rank_topk, torch.ops.aten.topk.default) + ) + ) + + exact_scatter = torch.export.export( + _ScatterOnly(), + ( + torch.zeros((1, 1, 262144), dtype=torch.float32), + torch.arange(4096, dtype=torch.int64).reshape(1, 1, 4096), + torch.ones((1, 1, 4096), dtype=torch.float32), + ), + strict=False, + ) + self.assertTrue( + _is_official_qat_unique_scatter( + _find_call(exact_scatter, torch.ops.aten.scatter.src) + ) + ) + wrong_rank_scatter = torch.export.export( + _ScatterOnly(), + ( + torch.zeros((1, 262144), dtype=torch.float32), + torch.arange(4096, dtype=torch.int64).reshape(1, 4096), + torch.ones((1, 4096), dtype=torch.float32), + ), + strict=False, + ) + self.assertFalse( + _is_official_qat_unique_scatter( + _find_call(wrong_rank_scatter, torch.ops.aten.scatter.src) + ) + ) diff --git a/examples/models/gemma4/tests/test_export_assistant_webgpu_artifacts.py b/examples/models/gemma4/tests/test_export_assistant_webgpu_artifacts.py new file mode 100644 index 00000000000..4ed46584241 --- /dev/null +++ b/examples/models/gemma4/tests/test_export_assistant_webgpu_artifacts.py @@ -0,0 +1,1058 @@ +# 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. + +# pyre-strict + +from __future__ import annotations + +import dataclasses +import json +import re +import tempfile +import unittest +from pathlib import Path +from typing import Any +from unittest import mock + +import executorch.backends.vulkan.custom_ops_lib # noqa: F401 +import executorch.extension.llm.custom_ops.custom_ops # noqa: F401 + +import torch + +from executorch.backends.vulkan.patterns.rope_hf import HfRotaryEmbeddingSinglePattern + +from executorch.examples.models.gemma4.eagle_webgpu_round import ( + OFFICIAL_QAT_CENTROID_TOP_K, + OFFICIAL_QAT_NUM_CENTROIDS, + OFFICIAL_QAT_SELECTED_TOKEN_COUNT, + OFFICIAL_QAT_TOKENS_PER_CENTROID, + select_qat_centroids, + validate_qat_token_ordering, + validate_selected_destinations, +) +from executorch.examples.models.gemma4.export_assistant_webgpu_artifacts import ( + _StaticAssistantQueryRopeLayer, + _UnusedAssistantRotaryEmbedding, + adapt_assistant_model_for_webgpu, + adapt_masked_embedding_for_webgpu, + QAT_VALIDATION_DONOR_SEQUENCE, + StaticAssistantMasks, + StaticAssistantQueryRope, + StaticAssistantSharedKVAttention, + UnfoldedAssistant, + validate_qat_centroid_scores, + validate_qat_selection_contract, +) +from executorch.examples.models.gemma4.webgpu_artifact_manifest import ( + ASSISTANT_CHECKPOINT_ACQUISITION, + ASSISTANT_MODEL_CONTRACT, +) + + +_VOCAB_SIZE = OFFICIAL_QAT_NUM_CENTROIDS * OFFICIAL_QAT_TOKENS_PER_CENTROID +_SLIDING_WINDOW = 512 +_HEX64 = re.compile(r"\A[0-9a-f]{64}\Z") +_HEX40 = re.compile(r"\A[0-9a-f]{40}\Z") +_HUB_REPOSITORY = re.compile(r"\A[A-Za-z0-9][A-Za-z0-9._-]*/[A-Za-z0-9._-]+\Z") +_FORBIDDEN_IDENTITY_MARKERS = ( + "/data/", # oss-closure-fixture + "/home/", # oss-closure-fixture + "/mnt/", # oss-closure-fixture + "fbsource", + "manifold", # oss-closure-fixture + "everstore", + "://", + "~", + "..", + " ", +) + + +class AssistantIdentityTest(unittest.TestCase): + """The pinned assistant checkpoint is a public hub coordinate plus digests.""" + + def test_repository_and_revision_are_pinned(self) -> None: + self.assertEqual( + ASSISTANT_CHECKPOINT_ACQUISITION["repo_id"], + "google/gemma-4-E2B-it-qat-q4_0-unquantized-assistant", + ) + self.assertEqual( + ASSISTANT_CHECKPOINT_ACQUISITION["revision"], + "ebc7e1a211354561464cb82ed6d886792138dcb6", + ) + self.assertRegex(str(ASSISTANT_CHECKPOINT_ACQUISITION["revision"]), _HEX40) + self.assertRegex( + str(ASSISTANT_CHECKPOINT_ACQUISITION["repo_id"]), _HUB_REPOSITORY + ) + + def test_every_pinned_digest_is_a_sha256(self) -> None: + files = ASSISTANT_CHECKPOINT_ACQUISITION["files"] + self.assertIsInstance(files, dict) + assert isinstance(files, dict) + self.assertEqual(set(files), {"config.json", "model.safetensors"}) + digests = set() + for name, identity in files.items(): + with self.subTest(name=name): + self.assertIsInstance(identity, dict) + assert isinstance(identity, dict) + self.assertRegex(str(identity["sha256"]), _HEX64) + self.assertGreater(int(identity["bytes"]), 0) + digests.add(identity["sha256"]) + self.assertEqual(len(digests), 2) + + def test_no_internal_path_or_receipt_is_pinned(self) -> None: + value = json.dumps( + ASSISTANT_CHECKPOINT_ACQUISITION, + sort_keys=True, + separators=(",", ":"), + ).lower() + for marker in _FORBIDDEN_IDENTITY_MARKERS: + with self.subTest(marker=marker): + self.assertNotIn(marker, value) + + def test_model_contract_is_the_official_assistant_shape(self) -> None: + self.assertEqual( + ASSISTANT_MODEL_CONTRACT, + { + "architecture": "Gemma4AssistantForCausalLM", + "backboneHiddenSize": 1536, + "hiddenSize": 256, + "modelType": "gemma4_assistant", + "numHiddenLayers": 4, + "vocabSize": 262144, + }, + ) + + def test_validation_donor_sequence_brackets_the_sliding_window(self) -> None: + self.assertEqual( + QAT_VALIDATION_DONOR_SEQUENCE, + (2, 16, 511, 512, 513, 514, 1024, 8960, 2), + ) + self.assertEqual(QAT_VALIDATION_DONOR_SEQUENCE[0], 2) + self.assertEqual(QAT_VALIDATION_DONOR_SEQUENCE[-1], 2) + for boundary in (_SLIDING_WINDOW - 1, _SLIDING_WINDOW, _SLIDING_WINDOW + 1): + self.assertIn(boundary, QAT_VALIDATION_DONOR_SEQUENCE) + + +class QatSelectionDimensionsTest(unittest.TestCase): + """2048 centroids x 128 tokens, 32 selected centroids, 4096 selected tokens.""" + + def test_official_selection_dimensions(self) -> None: + self.assertEqual(OFFICIAL_QAT_NUM_CENTROIDS, 2048) + self.assertEqual(OFFICIAL_QAT_TOKENS_PER_CENTROID, 128) + self.assertEqual(OFFICIAL_QAT_CENTROID_TOP_K, 32) + self.assertEqual(OFFICIAL_QAT_SELECTED_TOKEN_COUNT, 4096) + self.assertEqual( + OFFICIAL_QAT_SELECTED_TOKEN_COUNT, + OFFICIAL_QAT_CENTROID_TOP_K * OFFICIAL_QAT_TOKENS_PER_CENTROID, + ) + self.assertEqual(_VOCAB_SIZE, ASSISTANT_MODEL_CONTRACT["vocabSize"]) + + +def _token_ordering(seed: int = 0xE4A6) -> torch.Tensor: + generator = torch.Generator().manual_seed(seed) + return torch.randperm(_VOCAB_SIZE, generator=generator, dtype=torch.int64) + + +class QatTokenOrderingTest(unittest.TestCase): + def test_raw_and_logical_orderings_produce_the_same_evidence(self) -> None: + raw = _token_ordering() + logical = raw.reshape( + OFFICIAL_QAT_NUM_CENTROIDS, OFFICIAL_QAT_TOKENS_PER_CENTROID + ) + raw_evidence = validate_qat_token_ordering(raw) + logical_evidence = validate_qat_token_ordering(logical) + self.assertEqual(raw_evidence["rawShape"], [_VOCAB_SIZE]) + self.assertEqual(logical_evidence["rawShape"], [2048, 128]) + self.assertEqual(raw_evidence["shape"], [2048, 128]) + self.assertEqual(logical_evidence["shape"], [2048, 128]) + for key in ("max", "min", "numel", "sha256", "uniqueCount", "permutationExact"): + with self.subTest(key=key): + self.assertEqual(raw_evidence[key], logical_evidence[key]) + self.assertEqual(raw_evidence["min"], 0) + self.assertEqual(raw_evidence["max"], _VOCAB_SIZE - 1) + self.assertEqual(raw_evidence["numel"], _VOCAB_SIZE) + self.assertEqual(raw_evidence["uniqueCount"], _VOCAB_SIZE) + self.assertIs(raw_evidence["permutationExact"], True) + + def test_a_single_duplicated_entry_is_rejected(self) -> None: + duplicate = _token_ordering() + duplicate[-1] = duplicate[0] + with self.assertRaisesRegex(ValueError, "exact permutation"): + validate_qat_token_ordering(duplicate) + + def test_shifted_ordering_outside_the_vocabulary_is_rejected(self) -> None: + shifted = _token_ordering() + 1 + with self.assertRaisesRegex(ValueError, "exact permutation"): + validate_qat_token_ordering(shifted) + + def test_wrong_element_count_is_rejected(self) -> None: + with self.assertRaisesRegex(ValueError, "262144 entries"): + validate_qat_token_ordering(torch.arange(_VOCAB_SIZE - 1)) + + def test_non_integer_dtype_is_rejected(self) -> None: + with self.assertRaisesRegex(ValueError, "integer dtype"): + validate_qat_token_ordering(torch.arange(_VOCAB_SIZE, dtype=torch.float32)) + + +class QatSelectedDestinationTest(unittest.TestCase): + def test_thirty_two_centroids_map_to_4096_distinct_destinations(self) -> None: + ordering = _token_ordering() + scores = torch.arange(OFFICIAL_QAT_NUM_CENTROIDS, dtype=torch.float32).reshape( + 1, 1, OFFICIAL_QAT_NUM_CENTROIDS + ) + selected = select_qat_centroids(scores) + self.assertEqual(tuple(selected.shape), (OFFICIAL_QAT_CENTROID_TOP_K,)) + highest = OFFICIAL_QAT_NUM_CENTROIDS - 1 + self.assertEqual( + selected.tolist(), + list(range(highest, highest - OFFICIAL_QAT_CENTROID_TOP_K, -1)), + ) + destinations = validate_selected_destinations(ordering, selected) + self.assertEqual(destinations.numel(), OFFICIAL_QAT_SELECTED_TOKEN_COUNT) + self.assertEqual(torch.unique(destinations).numel(), destinations.numel()) + logical = ordering.reshape( + OFFICIAL_QAT_NUM_CENTROIDS, OFFICIAL_QAT_TOKENS_PER_CENTROID + ) + self.assertTrue( + torch.equal(destinations, logical[selected].reshape(-1)), + ) + + def test_duplicate_selected_centroids_are_rejected(self) -> None: + ordering = _token_ordering() + selected = torch.arange(OFFICIAL_QAT_CENTROID_TOP_K, dtype=torch.int64) + selected[-1] = selected[0] + with self.assertRaisesRegex(ValueError, "must be distinct"): + validate_selected_destinations(ordering, selected) + + def test_out_of_range_and_miscounted_centroids_are_rejected(self) -> None: + ordering = _token_ordering() + too_large = torch.arange(OFFICIAL_QAT_CENTROID_TOP_K, dtype=torch.int64) + too_large[0] = OFFICIAL_QAT_NUM_CENTROIDS + with self.assertRaisesRegex(ValueError, "index out of range"): + validate_selected_destinations(ordering, too_large) + negative = torch.arange(OFFICIAL_QAT_CENTROID_TOP_K, dtype=torch.int64) + negative[0] = -1 + with self.assertRaisesRegex(ValueError, "index out of range"): + validate_selected_destinations(ordering, negative) + with self.assertRaisesRegex(ValueError, "32 centroids"): + validate_selected_destinations( + ordering, torch.arange(31, dtype=torch.int64) + ) + + def test_non_permutation_ordering_is_rejected_before_selection(self) -> None: + duplicate = _token_ordering() + duplicate[-1] = duplicate[0] + with self.assertRaisesRegex(ValueError, "exact permutation"): + validate_selected_destinations( + duplicate, torch.arange(OFFICIAL_QAT_CENTROID_TOP_K, dtype=torch.int64) + ) + + def test_centroid_score_count_is_enforced(self) -> None: + with self.assertRaisesRegex(ValueError, "2048 entries"): + select_qat_centroids(torch.zeros(2047, dtype=torch.float32)) + + +class QatCentroidScoreTest(unittest.TestCase): + def _scores(self) -> torch.Tensor: + return torch.arange(OFFICIAL_QAT_NUM_CENTROIDS, dtype=torch.float32).reshape( + 1, 1, OFFICIAL_QAT_NUM_CENTROIDS + ) + + def test_strictly_ordered_scores_produce_a_positive_boundary_gap(self) -> None: + evidence = validate_qat_centroid_scores(self._scores()) + self.assertIs(evidence["allFinite"], True) + self.assertIs(evidence["top32PairwiseDistinct"], True) + self.assertEqual(evidence["boundaryGap"], 1.0) + for key in ( + "indicesSha256", + "top33IndicesSha256", + "top33ValuesSha256", + "valuesSha256", + ): + with self.subTest(key=key): + self.assertRegex(str(evidence[key]), _HEX64) + + def test_a_tie_at_the_32_33_boundary_is_rejected(self) -> None: + scores = self._scores() + scores[..., OFFICIAL_QAT_NUM_CENTROIDS - 33] = scores[ + ..., OFFICIAL_QAT_NUM_CENTROIDS - 32 + ] + with self.assertRaisesRegex(ValueError, "boundary gap"): + validate_qat_centroid_scores(scores) + + def test_a_tie_inside_the_top_32_is_rejected(self) -> None: + scores = self._scores() + scores[..., OFFICIAL_QAT_NUM_CENTROIDS - 2] = scores[ + ..., OFFICIAL_QAT_NUM_CENTROIDS - 1 + ] + with self.assertRaisesRegex(ValueError, "pairwise distinct"): + validate_qat_centroid_scores(scores) + + def test_non_finite_scores_are_rejected(self) -> None: + for bad in (float("nan"), float("inf"), float("-inf")): + with self.subTest(bad=bad): + scores = self._scores() + scores[..., 0] = bad + with self.assertRaisesRegex(ValueError, "finite fp32"): + validate_qat_centroid_scores(scores) + + def test_shape_and_dtype_are_pinned(self) -> None: + with self.assertRaisesRegex(ValueError, r"shape \[1, 1, 2048\]"): + validate_qat_centroid_scores( + torch.zeros((1, OFFICIAL_QAT_NUM_CENTROIDS), dtype=torch.float32) + ) + with self.assertRaisesRegex(ValueError, "finite fp32"): + validate_qat_centroid_scores( + torch.zeros((1, 1, OFFICIAL_QAT_NUM_CENTROIDS), dtype=torch.float64) + ) + + +class _MaskedEmbeddingFixture(torch.nn.Module): + """Minimal stand-in for the QAT masked embedding head.""" + + def __init__(self, hidden_size: int = 4, *, with_lm_embed: bool = True) -> None: + super().__init__() + generator = torch.Generator().manual_seed(11) + self.centroids = torch.nn.Linear( + hidden_size, OFFICIAL_QAT_NUM_CENTROIDS, bias=False + ) + with torch.no_grad(): + self.centroids.weight.copy_( + torch.randn( + (OFFICIAL_QAT_NUM_CENTROIDS, hidden_size), generator=generator + ) + ) + self.register_buffer("token_ordering", _token_ordering(), persistent=False) + if with_lm_embed: + embed = torch.nn.Embedding(_VOCAB_SIZE, hidden_size) + with torch.no_grad(): + embed.weight.copy_( + torch.randn((_VOCAB_SIZE, hidden_size), generator=generator) + ) + self._lm_embed = embed + + def forward( + self, hidden_states: torch.Tensor, lm_head_weight: torch.Tensor + ) -> torch.Tensor: + del lm_head_weight + return self.centroids(hidden_states) + + +class MaskedEmbeddingAdaptationTest(unittest.TestCase): + """The static head must reproduce a masked full-vocabulary logit row.""" + + def test_static_head_matches_the_dense_masked_reference(self) -> None: + head = _MaskedEmbeddingFixture() + ordering = head.token_ordering.clone() + adapt_masked_embedding_for_webgpu(head) + generator = torch.Generator().manual_seed(5) + hidden = torch.randn((1, 1, 4), generator=generator) + + produced = head(hidden, torch.empty(0)) + centroid_logits = head.centroids(hidden) + top_k = torch.topk(centroid_logits, OFFICIAL_QAT_CENTROID_TOP_K, dim=-1).indices + selected = ( + ordering.reshape( + OFFICIAL_QAT_NUM_CENTROIDS, OFFICIAL_QAT_TOKENS_PER_CENTROID + )[top_k.reshape(-1)] + .reshape(1, 1, OFFICIAL_QAT_SELECTED_TOKEN_COUNT) + .to(torch.int64) + ) + dense = torch.nn.functional.linear(hidden, head._lm_embed.weight) + reference = torch.full_like(dense, torch.finfo(torch.float32).min) + reference.scatter_(-1, selected, dense.gather(-1, selected)) + + self.assertEqual(tuple(produced.shape), (1, 1, _VOCAB_SIZE)) + self.assertEqual(produced.dtype, torch.float32) + torch.testing.assert_close(produced, reference, atol=1e-5, rtol=1e-5) + + def test_unselected_tokens_stay_masked_out(self) -> None: + head = _MaskedEmbeddingFixture() + adapt_masked_embedding_for_webgpu(head) + generator = torch.Generator().manual_seed(6) + produced = head(torch.randn((1, 1, 4), generator=generator), torch.empty(0)) + masked = produced == torch.finfo(torch.float32).min + self.assertEqual(int((~masked).sum().item()), OFFICIAL_QAT_SELECTED_TOKEN_COUNT) + + def test_adaptation_publishes_the_static_selection_buffers(self) -> None: + head = _MaskedEmbeddingFixture() + ordering = head.token_ordering.clone() + adapt_masked_embedding_for_webgpu(head) + self.assertEqual( + tuple(head._webgpu_token_ordering.shape), + (OFFICIAL_QAT_NUM_CENTROIDS, OFFICIAL_QAT_TOKENS_PER_CENTROID), + ) + self.assertEqual(head._webgpu_token_ordering.dtype, torch.float32) + self.assertTrue( + torch.equal( + head._webgpu_token_ordering.to(torch.int64).reshape(-1), ordering + ) + ) + self.assertEqual(tuple(head._webgpu_output_template.shape), (1, 1, _VOCAB_SIZE)) + self.assertTrue( + torch.all(head._webgpu_output_template == torch.finfo(torch.float32).min) + ) + + def test_adaptation_requires_a_quantized_lm_head(self) -> None: + with self.assertRaisesRegex(ValueError, "quantized LM head"): + adapt_masked_embedding_for_webgpu( + _MaskedEmbeddingFixture(with_lm_embed=False) + ) + + +class StaticAssistantMaskTest(unittest.TestCase): + """The sliding mask exposes the newest `sliding_window + 1` donor rows.""" + + def test_masks_track_the_donor_window(self) -> None: + masks = StaticAssistantMasks(max_seq_len=8960) + blocked = torch.finfo(torch.float32).min + for donor_length in (2, 511, 512, 513, 514, 8960): + with self.subTest(donor_length=donor_length): + donor = torch.zeros((1, 1, donor_length, 1)) + full, sliding = masks(donor, donor) + self.assertEqual(tuple(full.shape), (1, donor_length)) + self.assertEqual(tuple(sliding.shape), (1, donor_length)) + self.assertTrue(torch.equal(full, torch.zeros_like(full))) + visible = min(donor_length, _SLIDING_WINDOW + 1) + self.assertTrue( + torch.equal(sliding[:, -visible:], torch.zeros((1, visible))) + ) + if donor_length > visible: + self.assertTrue( + torch.equal( + sliding[:, :-visible], + torch.full((1, donor_length - visible), blocked), + ) + ) + + def test_full_and_sliding_donors_may_have_different_lengths(self) -> None: + masks = StaticAssistantMasks(max_seq_len=4096) + full, sliding = masks(torch.zeros((1, 1, 4096, 1)), torch.zeros((1, 1, 600, 1))) + self.assertEqual(tuple(full.shape), (1, 4096)) + self.assertEqual(tuple(sliding.shape), (1, 600)) + + def test_degenerate_capacities_are_rejected(self) -> None: + with self.assertRaisesRegex(ValueError, "attention-mask capacity"): + StaticAssistantMasks(max_seq_len=1) + with self.assertRaisesRegex(ValueError, "attention-mask capacity"): + StaticAssistantMasks(max_seq_len=1024, sliding_window=0) + + +class _AssistantRotaryEmbedding(torch.nn.Module): + """Per-layer-type rotary table with the HF rotate-half layout.""" + + def __init__(self, head_dim: int) -> None: + super().__init__() + half = head_dim // 2 + exponent = torch.arange(0, half, dtype=torch.float32) / half + self.register_buffer( + "full_inv_freq", 1.0 / (10000.0**exponent), persistent=False + ) + self.register_buffer( + "sliding_inv_freq", 1.0 / (100.0**exponent), persistent=False + ) + + def forward( + self, + x: torch.Tensor, + position_ids: torch.Tensor, + layer_type: str | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + inv_freq = ( + self.sliding_inv_freq + if layer_type == "sliding_attention" + else self.full_inv_freq + ) + angles = position_ids.to(torch.float32).unsqueeze(-1) * inv_freq + table = torch.cat((angles, angles), dim=-1) + return table.cos().to(x.dtype), table.sin().to(x.dtype) + + +class _ReferenceSharedKVAttention(torch.nn.Module): + """Unadapted layer: rotary tables recomputed per call, plain fp32 SDPA.""" + + def __init__( + self, + rotary: _AssistantRotaryEmbedding, + *, + layer_type: str, + hidden_size: int = 8, + num_attention_heads: int = 2, + head_dim: int = 4, + is_kv_shared_layer: bool = True, + ) -> None: + super().__init__() + self.is_kv_shared_layer = is_kv_shared_layer + self.layer_type = layer_type + self.head_dim = head_dim + self.num_attention_heads = num_attention_heads + self.q_proj = torch.nn.Linear( + hidden_size, num_attention_heads * head_dim, bias=False + ) + self.q_norm = torch.nn.LayerNorm(head_dim) + self.o_proj = torch.nn.Linear( + num_attention_heads * head_dim, hidden_size, bias=False + ) + self.rotary = rotary + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor, + shared_kv_states: dict[str, tuple[torch.Tensor, torch.Tensor]], + position_ids: torch.Tensor, + ) -> torch.Tensor: + input_shape = hidden_states.shape[:-1] + query = self.q_proj(hidden_states).view( + *input_shape, self.num_attention_heads, self.head_dim + ) + query = self.q_norm(query) + cos, sin = self.rotary(query, position_ids, self.layer_type) + query = HfRotaryEmbeddingSinglePattern()(query, cos.squeeze(0), sin.squeeze(0)) + key, value = shared_kv_states[self.layer_type] + attention = torch.nn.functional.scaled_dot_product_attention( + query.transpose(1, 2), key, value, attn_mask=attention_mask, scale=1.0 + ).transpose(1, 2) + return self.o_proj(attention.reshape(*input_shape, -1)) + + +def _donor_states( + donor_length: int, heads: int = 2, head_dim: int = 4, seed: int = 3 +) -> dict[str, tuple[torch.Tensor, torch.Tensor]]: + generator = torch.Generator().manual_seed(seed) + states: dict[str, tuple[torch.Tensor, torch.Tensor]] = {} + for layer_type in ("full_attention", "sliding_attention"): + states[layer_type] = ( + torch.randn((1, heads, donor_length, head_dim), generator=generator), + torch.randn((1, heads, donor_length, head_dim), generator=generator), + ) + return states + + +class AdaptedAssistantAttentionTest(unittest.TestCase): + """`StaticAssistant*` must reproduce the unadapted eager attention output.""" + + MAX_SEQ_LEN = 1024 + TOLERANCE = 1e-4 + + def _adapted( + self, reference: _ReferenceSharedKVAttention, rotary: _AssistantRotaryEmbedding + ) -> StaticAssistantSharedKVAttention: + tables = StaticAssistantQueryRope(rotary, max_seq_len=self.MAX_SEQ_LEN) + source = getattr(tables, reference.layer_type) + return StaticAssistantSharedKVAttention( + reference, + _StaticAssistantQueryRopeLayer(source.freqs_cos, source.freqs_sin), + ) + + def test_precomputed_rope_tables_match_the_live_rotary_module(self) -> None: + torch.manual_seed(0) + rotary = _AssistantRotaryEmbedding(head_dim=4) + tables = StaticAssistantQueryRope(rotary, max_seq_len=self.MAX_SEQ_LEN) + query = torch.randn((1, 1, 2, 4)) + for layer_type in ("full_attention", "sliding_attention"): + for position in (0, 1, 511, 512, 513, self.MAX_SEQ_LEN - 1): + with self.subTest(layer_type=layer_type, position=position): + position_ids = torch.tensor([[position]], dtype=torch.int64) + cos, sin = rotary(query, position_ids, layer_type) + reference = HfRotaryEmbeddingSinglePattern()( + query, cos.squeeze(0), sin.squeeze(0) + ) + source = getattr(tables, layer_type) + produced = _StaticAssistantQueryRopeLayer( + source.freqs_cos, source.freqs_sin + )(query, position_ids) + torch.testing.assert_close( + produced, reference, atol=self.TOLERANCE, rtol=self.TOLERANCE + ) + + def test_rope_tables_are_layer_type_specific(self) -> None: + torch.manual_seed(0) + tables = StaticAssistantQueryRope( + _AssistantRotaryEmbedding(head_dim=4), max_seq_len=self.MAX_SEQ_LEN + ) + self.assertFalse( + torch.equal( + tables.full_attention.freqs_cos, tables.sliding_attention.freqs_cos + ) + ) + self.assertEqual( + tuple(tables.full_attention.freqs_cos.shape), (self.MAX_SEQ_LEN, 4) + ) + + def test_adapted_layer_matches_the_unadapted_layer(self) -> None: + torch.manual_seed(0) + rotary = _AssistantRotaryEmbedding(head_dim=4) + masks = StaticAssistantMasks(self.MAX_SEQ_LEN) + hidden = torch.randn((1, 1, 8)) + for layer_type in ("full_attention", "sliding_attention"): + for donor_length in (2, 511, 512, 513, 514, self.MAX_SEQ_LEN): + with self.subTest(layer_type=layer_type, donor_length=donor_length): + reference = _ReferenceSharedKVAttention( + rotary, layer_type=layer_type + ) + adapted = self._adapted(reference, rotary) + shared = _donor_states(donor_length) + full_mask, sliding_mask = masks( + shared["full_attention"][0], shared["sliding_attention"][0] + ) + mask = full_mask if layer_type == "full_attention" else sliding_mask + position_ids = torch.tensor([[donor_length - 1]], dtype=torch.int64) + expected = reference(hidden, mask, shared, position_ids) + produced, extra = adapted(hidden, None, mask, shared, position_ids) + self.assertIsNone(extra) + self.assertEqual(produced.shape, expected.shape) + torch.testing.assert_close( + produced, expected, atol=self.TOLERANCE, rtol=self.TOLERANCE + ) + + def test_adapted_layer_reads_only_its_own_shared_kv_entry(self) -> None: + torch.manual_seed(0) + rotary = _AssistantRotaryEmbedding(head_dim=4) + reference = _ReferenceSharedKVAttention(rotary, layer_type="full_attention") + adapted = self._adapted(reference, rotary) + hidden = torch.randn((1, 1, 8)) + mask = torch.zeros((1, 8)) + position_ids = torch.tensor([[7]], dtype=torch.int64) + shared = _donor_states(8) + baseline, _ = adapted(hidden, None, mask, shared, position_ids) + swapped = dict(shared) + swapped["sliding_attention"] = _donor_states(8, seed=99)["sliding_attention"] + produced, _ = adapted(hidden, None, mask, swapped, position_ids) + torch.testing.assert_close( + produced, baseline, atol=self.TOLERANCE, rtol=self.TOLERANCE + ) + + def test_a_layer_without_shared_target_kv_is_rejected(self) -> None: + torch.manual_seed(0) + rotary = _AssistantRotaryEmbedding(head_dim=4) + reference = _ReferenceSharedKVAttention( + rotary, layer_type="full_attention", is_kv_shared_layer=False + ) + tables = StaticAssistantQueryRope(rotary, max_seq_len=self.MAX_SEQ_LEN) + with self.assertRaisesRegex(ValueError, "shared target KV"): + StaticAssistantSharedKVAttention(reference, tables.full_attention) + + def test_a_query_beyond_the_rope_table_is_rejected(self) -> None: + torch.manual_seed(0) + rotary = _AssistantRotaryEmbedding(head_dim=4) + tables = StaticAssistantQueryRope(rotary, max_seq_len=self.MAX_SEQ_LEN) + layer = _StaticAssistantQueryRopeLayer( + tables.full_attention.freqs_cos, tables.full_attention.freqs_sin + ) + query = torch.randn((1, 1, 2, 4)) + with self.assertRaises(RuntimeError): + layer(query, torch.tensor([[self.MAX_SEQ_LEN]], dtype=torch.int64)) + with self.assertRaises(RuntimeError): + layer(query, torch.tensor([[-1]], dtype=torch.int64)) + + +@dataclasses.dataclass +class _AssistantOutput: + logits: torch.Tensor + last_hidden_state: torch.Tensor + + +class _RoutingAssistant(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.calls: list[dict[str, Any]] = [] + + def forward( + self, + inputs_embeds: torch.Tensor, + attention_mask: torch.Tensor | None, + position_ids: torch.Tensor, + shared_kv_states: dict[str, tuple[torch.Tensor, torch.Tensor]], + use_cache: bool, + ) -> _AssistantOutput: + self.calls.append( + { + "attention_mask": attention_mask, + "inputs_embeds": inputs_embeds, + "position_ids": position_ids, + "shared_kv_states": shared_kv_states, + "use_cache": use_cache, + } + ) + return _AssistantOutput( + logits=torch.full((1, 1, 3), 7.0), + last_hidden_state=torch.full((1, 1, 2), 9.0), + ) + + +class UnfoldedAssistantRoutingTest(unittest.TestCase): + def test_donor_tensors_are_routed_to_their_layer_types(self) -> None: + inner = _RoutingAssistant() + wrapper = UnfoldedAssistant(inner) + embeds = torch.zeros((1, 1, 4)) + position_ids = torch.tensor([[5]], dtype=torch.int64) + full_k, full_v, sliding_k, sliding_v = ( + torch.full((1, 2, 3, 4), float(marker)) for marker in range(4) + ) + + logits, hidden = wrapper( + embeds, position_ids, full_k, full_v, sliding_k, sliding_v + ) + + self.assertEqual(len(inner.calls), 1) + call = inner.calls[0] + self.assertIsNone(call["attention_mask"]) + self.assertIs(call["use_cache"], False) + self.assertIs(call["inputs_embeds"], embeds) + self.assertIs(call["position_ids"], position_ids) + self.assertEqual( + sorted(call["shared_kv_states"]), ["full_attention", "sliding_attention"] + ) + self.assertIs(call["shared_kv_states"]["full_attention"][0], full_k) + self.assertIs(call["shared_kv_states"]["full_attention"][1], full_v) + self.assertIs(call["shared_kv_states"]["sliding_attention"][0], sliding_k) + self.assertIs(call["shared_kv_states"]["sliding_attention"][1], sliding_v) + self.assertEqual(logits.tolist(), torch.full((1, 1, 3), 7.0).tolist()) + self.assertEqual(hidden.tolist(), torch.full((1, 1, 2), 9.0).tolist()) + + +class _AssistantLayer(torch.nn.Module): + def __init__(self, self_attn: _ReferenceSharedKVAttention) -> None: + super().__init__() + self.self_attn = self_attn + + +class _AssistantBackbone(torch.nn.Module): + def __init__( + self, rotary: _AssistantRotaryEmbedding, layer_types: list[str] + ) -> None: + super().__init__() + self.rotary_emb = rotary + self.layers = torch.nn.ModuleList( + [ + _AssistantLayer(_ReferenceSharedKVAttention(rotary, layer_type=name)) + for name in layer_types + ] + ) + + +class _AssistantModel(torch.nn.Module): + def __init__(self, layer_types: list[str] | None = None) -> None: + super().__init__() + types = layer_types or [ + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + ] + self.model = _AssistantBackbone(_AssistantRotaryEmbedding(head_dim=4), types) + self.masked_embedding = _MaskedEmbeddingFixture() + + +class AssistantAdaptationTest(unittest.TestCase): + MAX_SEQ_LEN = 1024 + + def test_adaptation_rewires_every_layer_and_the_masked_head(self) -> None: + torch.manual_seed(0) + assistant = _AssistantModel() + adapt_assistant_model_for_webgpu(assistant, max_seq_len=self.MAX_SEQ_LEN) + self.assertIsInstance( + assistant.model.rotary_emb, _UnusedAssistantRotaryEmbedding + ) + self.assertIsInstance(assistant._webgpu_static_masks, StaticAssistantMasks) + self.assertEqual(assistant._webgpu_static_masks.max_seq_len, self.MAX_SEQ_LEN) + self.assertEqual( + [layer.self_attn.layer_type for layer in assistant.model.layers], + [ + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + ], + ) + for layer in assistant.model.layers: + self.assertIsInstance(layer.self_attn, StaticAssistantSharedKVAttention) + self.assertTrue(hasattr(assistant.masked_embedding, "_webgpu_token_ordering")) + + def test_unexpected_layer_types_are_rejected(self) -> None: + torch.manual_seed(0) + sliding = "sliding_attention" + for layer_types in ( + ["full_attention", sliding, sliding, sliding], + [sliding, sliding, "full_attention"], + ): + with self.subTest(layer_types=layer_types): + with self.assertRaisesRegex(ValueError, "assistant layer types"): + adapt_assistant_model_for_webgpu( + _AssistantModel(layer_types), max_seq_len=self.MAX_SEQ_LEN + ) + + def test_static_masks_replace_the_caller_supplied_attention_mask(self) -> None: + torch.manual_seed(0) + assistant = _AssistantModel() + adapt_assistant_model_for_webgpu(assistant, max_seq_len=self.MAX_SEQ_LEN) + shared = _donor_states(600) + masks = assistant.create_attention_masks(torch.zeros((1, 1, 8)), None, shared) + self.assertEqual(sorted(masks), ["full_attention", "sliding_attention"]) + self.assertEqual(tuple(masks["full_attention"].shape), (1, 600)) + self.assertEqual(tuple(masks["sliding_attention"].shape), (1, 600)) + with self.assertRaisesRegex(ValueError, "attention_mask=None"): + assistant.create_attention_masks( + torch.zeros((1, 1, 8)), torch.zeros((1, 600)), shared + ) + + +class _QatTextConfig: + global_head_dim = 1 + head_dim = 1 + + +class _QatAssistantConfig: + backbone_hidden_size = 1 + + def get_text_config(self) -> _QatTextConfig: + return _QatTextConfig() + + +class _QatCentroidHead(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + generator = torch.Generator().manual_seed(17) + self.centroids = torch.nn.Linear(1, OFFICIAL_QAT_NUM_CENTROIDS, bias=False) + with torch.no_grad(): + self.centroids.weight.copy_( + torch.randn((OFFICIAL_QAT_NUM_CENTROIDS, 1), generator=generator) + ) + + def forward(self, hidden: torch.Tensor) -> torch.Tensor: + return self.centroids(hidden) + + +class _QatValidationAssistant(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.config = _QatAssistantConfig() + self.masked_embedding = _QatCentroidHead() + + +class _QatValidationWrapper(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.assistant = _QatValidationAssistant() + + def forward(self, *inputs: torch.Tensor) -> torch.Tensor: + return self.assistant.masked_embedding(inputs[0][..., :1]) + + +class QatSelectionContractTest(unittest.TestCase): + """The QAT receipt replays donor length 2 and brackets the window boundaries.""" + + def test_receipt_covers_the_documented_donor_axes(self) -> None: + evidence = validate_qat_selection_contract( + # pyre-ignore[6]: the wrapper duck-types `UnfoldedAssistant`. + _QatValidationWrapper(), + max_donor_len=8960, + ) + self.assertEqual( + evidence["donorSequence"], [2, 16, 511, 512, 513, 514, 1024, 8960, 2] + ) + self.assertEqual( + evidence["selectionContract"], + { + "centroidTopK": 32, + "numCentroids": 2048, + "selectedTokenCount": 4096, + "tokensPerCentroid": 128, + }, + ) + cases = evidence["cases"] + self.assertEqual(len(cases), 9) + self.assertEqual([case["caseIndex"] for case in cases], list(range(9))) + self.assertEqual(cases[0]["inputSha256"], cases[-1]["inputSha256"]) + self.assertEqual(cases[0]["topk"], cases[-1]["topk"]) + self.assertNotEqual(cases[0]["inputSha256"], cases[1]["inputSha256"]) + + def test_short_capacity_truncates_the_sequence_but_keeps_the_replay(self) -> None: + evidence = validate_qat_selection_contract( + # pyre-ignore[6]: the wrapper duck-types `UnfoldedAssistant`. + _QatValidationWrapper(), + max_donor_len=512, + ) + self.assertEqual(evidence["donorSequence"], [2, 16, 511, 512, 2]) + + def test_capacity_below_the_replay_length_is_rejected(self) -> None: + with self.assertRaisesRegex(ValueError, "replay at donor length 2"): + validate_qat_selection_contract( + # pyre-ignore[6]: the wrapper duck-types `UnfoldedAssistant`. + _QatValidationWrapper(), + max_donor_len=1, + ) + + +class _FakeExportGraphModule: + def __init__(self) -> None: + self.meta: dict[str, object] = { + "gemma4K2Abi": {"fixture": "k2_abi"}, + "gemma4QATSelectionEvidence": {"fixture": "qat_selection"}, + "gemma4TargetCheckpointEvidence": {"fixture": "target_checkpoint"}, + } + + +class _FakeK2Program: + def __init__(self) -> None: + self.graph_module = _FakeExportGraphModule() + + +class _FakeExecutorchProgram: + def __init__(self) -> None: + self._tensor_data = {f"constants_{index}": object() for index in range(3)} + + def write_to_file(self, output: Any) -> None: + output.write(b"pte") + + def write_tensor_data_to_file(self, directory: str) -> None: + root = Path(directory) + for tag in self._tensor_data: + (root / f"{tag}.ptd").write_bytes(tag.encode("utf-8")) + + +class SpeculativeExportPublicationTest(unittest.TestCase): + def setUp(self) -> None: + from executorch.examples.models.gemma4 import export_speculative + + self.export_module = export_speculative + self._temporary_directory = tempfile.TemporaryDirectory() + self.addCleanup(self._temporary_directory.cleanup) + self.root = Path(self._temporary_directory.name) + self.output_root = self.root / "sealed" + self.output = self.output_root / "model.pte" + self.receipt = self.root / "receipts" / "manifest.json" + self.source = self.root / "source" / "source.json" + self.source.parent.mkdir() + self.source.write_bytes(b"sealed-source-receipt") + + def _run_export(self, validation: Any = None) -> Path: + from executorch.examples.models.gemma4 import webgpu_artifact_manifest + + validator = validation if validation is not None else mock.Mock() + with ( + mock.patch.object( + self.export_module, + "validate_assistant_checkpoint", + return_value={"fixture": "assistant_checkpoint"}, + ), + mock.patch.object( + self.export_module, + "build_k2_round_program", + return_value=_FakeK2Program(), + ), + mock.patch.object( + self.export_module, + "_lower_k2_round", + return_value=( + _FakeExecutorchProgram(), + {"fixture": "lowering"}, + ), + ), + mock.patch.object( + webgpu_artifact_manifest, + "create_mtp_manifest", + return_value={"schema_version": 1}, + ), + mock.patch.object( + webgpu_artifact_manifest, + "validate_mtp_manifest", + side_effect=validator, + ), + ): + return self.export_module.export_speculative( + self.root / "target", + self.root / "assistant", + self.output, + self.receipt, + source_receipt_path=self.source, + ) + + def test_external_source_receipt_is_published_without_moving_the_input( + self, + ) -> None: + source_bytes = self.source.read_bytes() + + self.assertEqual(self.receipt, self._run_export()) + + self.assertEqual(source_bytes, self.source.read_bytes()) + self.assertEqual( + source_bytes, (self.output_root / self.source.name).read_bytes() + ) + self.assertTrue(self.output.is_file()) + self.assertTrue(self.receipt.is_file()) + + def test_real_export_path_delegates_to_atomic_finalizer(self) -> None: + finalizer = mock.Mock(return_value=self.receipt) + with mock.patch.object( + self.export_module, + "finalize_mtp_export", + finalizer, + ): + result = self._run_export() + + self.assertEqual(result, self.receipt) + finalizer.assert_called_once() + ( + staging, + output, + receipt, + staged_pte, + staged_ptds, + source_receipt, + evidence, + ) = finalizer.call_args.args + self.assertEqual(staging, staged_pte.parent) + self.assertEqual(output, self.output) + self.assertEqual(receipt, self.receipt) + self.assertEqual(staged_pte.name, self.output.name) + self.assertEqual( + [path.name for path in staged_ptds], + [f"constants_{index}.ptd" for index in range(3)], + ) + self.assertEqual(source_receipt, self.source) + self.assertEqual( + { + "assistant_checkpoint": {"fixture": "assistant_checkpoint"}, + "k2_abi": {"fixture": "k2_abi"}, + "lowering": {"fixture": "lowering"}, + "qat_selection": {"fixture": "qat_selection"}, + "target_checkpoint": {"fixture": "target_checkpoint"}, + }, + evidence, + ) + + def test_source_basename_alias_is_rejected_before_publication(self) -> None: + self.source = self.source.with_name(self.output.name) + self.source.write_bytes(b"alias") + + with self.assertRaisesRegex(ValueError, "duplicate normalized artifact path"): + self._run_export() + + self.assertFalse(self.output.exists()) + self.assertFalse(self.receipt.exists()) + self.assertEqual(b"alias", self.source.read_bytes()) + + def test_failed_final_validation_rolls_back_source_and_model_artifacts( + self, + ) -> None: + calls = 0 + + def fail_after_publication(_root: Path, _manifest: object) -> None: + nonlocal calls + calls += 1 + if calls == 2: + raise ValueError("injected final validation failure") + + with self.assertRaisesRegex(ValueError, "injected final validation failure"): + self._run_export(fail_after_publication) + + self.assertEqual(b"sealed-source-receipt", self.source.read_bytes()) + self.assertFalse(self.output.exists()) + self.assertFalse((self.output_root / self.source.name).exists()) + self.assertFalse(self.receipt.exists()) diff --git a/examples/models/gemma4/tests/test_export_partitioners.py b/examples/models/gemma4/tests/test_export_partitioners.py index 8e769ca9179..affb24b4c11 100644 --- a/examples/models/gemma4/tests/test_export_partitioners.py +++ b/examples/models/gemma4/tests/test_export_partitioners.py @@ -6,12 +6,17 @@ import unittest -from executorch.backends.vulkan.op_registry import vulkan_supported_ops +import torch + +from executorch.backends.vulkan.op_registry import has_impl, vulkan_supported_ops from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner from executorch.examples.models.gemma4.webgpu_partitioner import ( + _extra_op_features, _webgpu_allowlist, build_webgpu_partitioner, + Gemma4WebGPUPartitioner, ) +from executorch.exir import EdgeCompileConfig, to_edge from executorch.exir.dialects._ops import ops as exir_ops @@ -51,3 +56,667 @@ def test_restricted_allowlist_includes_symbolic_select(self) -> None: def test_emb8_fails_closed(self) -> None: with self.assertRaisesRegex(ValueError, "emb8"): build_webgpu_partitioner("8da4w+emb8") + + +_QUERY_SHAPE = (1, 4, 8, 256) +_KV_SHAPE = (1, 4, 1, 256) +_MASK_SHAPE = (4, 4) +_ROPE_SHAPE = (1, 4, 8, 256) +_FREQS_SHAPE = (4, 128) + + +class _ScopedLayer(torch.nn.Module): + pass + + +class _GraphProgram: + """Stands in for the ExportedProgram surface the scoped rewriters read.""" + + __slots__ = ("graph_module",) + + def __init__(self, graph_module: torch.fx.GraphModule) -> None: + self.graph_module = graph_module + + +def _meta(shape: tuple[int, ...], dtype: torch.dtype = torch.float32) -> torch.Tensor: + return torch.empty(shape, dtype=dtype, device="meta") + + +def _placeholder( + graph: torch.fx.Graph, + name: str, + shape: tuple[int, ...], + dtype: torch.dtype = torch.float32, +) -> torch.fx.Node: + node = graph.placeholder(name) + node.meta["val"] = _meta(shape, dtype) + return node + + +# Each layout defeats a different stand-in for scope matching: `interleaved` a +# positional slice, `nested` an equality/prefix test on the module path. +_SCOPE_LAYOUTS = ("blocks", "interleaved", "nested") + + +def _scope_meta(scope: str, index: int, layout: str) -> dict[str, tuple[str, type]]: + path = ( + f"decoder.{scope}_layers.{index}" if layout == "nested" else f"{scope}.{index}" + ) + return {f"L__self___{scope}_{index}": (path, _ScopedLayer)} + + +def _scope_plan(target: int, assistant: int, unscoped: int, layout: str) -> list[str]: + queues = [["target"] * target, ["assistant"] * assistant, [""] * unscoped] + if layout != "interleaved": + return [scope for queue in queues for scope in queue] + plan: list[str] = [] + while any(queues): + for queue in queues: + if queue: + plan.append(queue.pop()) + return plan + + +def _module_paths(node: torch.fx.Node) -> list[str]: + stack = node.meta.get("nn_module_stack") or {} + return [entry[0] for entry in stack.values()] + + +def _nodes_with_target(graph: torch.fx.Graph, target: object) -> list[torch.fx.Node]: + return [node for node in graph.nodes if node.target == target] + + +def _scoped_graph_module( + graph: torch.fx.Graph, + target_op: object, + call_args: tuple[object, ...], + result_shape: tuple[int, ...], + plan: list[str], + layout: str, +) -> torch.fx.GraphModule: + emitted: dict[str, int] = {} + calls: list[torch.fx.Node] = [] + for scope in plan: + index = emitted.get(scope, 0) + emitted[scope] = index + 1 + node = graph.call_function(target_op, call_args) + node.meta["val"] = _meta(result_shape) + if scope: + node.meta["nn_module_stack"] = _scope_meta(scope, index, layout) + calls.append(node) + graph.output(tuple(calls)) + graph.lint() + return torch.fx.GraphModule(torch.nn.Module(), graph) + + +def _sdpa_program( + target: int, assistant: int, unscoped: int, layout: str = "blocks" +) -> _GraphProgram: + graph = torch.fx.Graph() + query = _placeholder(graph, "query", _QUERY_SHAPE) + key = _placeholder(graph, "key", _KV_SHAPE) + value = _placeholder(graph, "value", _KV_SHAPE) + mask = _placeholder(graph, "mask", _MASK_SHAPE) + return _GraphProgram( + _scoped_graph_module( + graph, + exir_ops.edge.llama.custom_sdpa.default, + (query, key, value, 0, mask, 0.0, False, 1.0), + _QUERY_SHAPE, + _scope_plan(target, assistant, unscoped, layout), + layout, + ) + ) + + +def _rope_program( + target: int, assistant: int, unscoped: int, layout: str = "blocks" +) -> _GraphProgram: + graph = torch.fx.Graph() + activations = _placeholder(graph, "x", _ROPE_SHAPE) + freqs_cos = _placeholder(graph, "freqs_cos", _FREQS_SHAPE) + freqs_sin = _placeholder(graph, "freqs_sin", _FREQS_SHAPE) + return _GraphProgram( + _scoped_graph_module( + graph, + exir_ops.edge.et_vk.apply_rotary_emb_hf_single.default, + (activations, freqs_cos, freqs_sin, 0), + _ROPE_SHAPE, + _scope_plan(target, assistant, unscoped, layout), + layout, + ) + ) + + +class _MtpPartitionerTest(unittest.TestCase): + def setUp(self) -> None: + # Deferred: the plain cases above must load without D8 or the custom-ops lib. + import executorch.examples.models.gemma4.webgpu_partitioner as mtp + import executorch.extension.llm.custom_ops.custom_ops # noqa: F401 + + self.mtp = mtp + + +class MTPScopedSDPATest(_MtpPartitionerTest): + def test_exactly_the_target_scoped_sdpa_sites_are_rewritten(self) -> None: + custom = exir_ops.edge.llama.custom_sdpa.default + for layout in _SCOPE_LAYOUTS: + with self.subTest(layout=layout): + program = _sdpa_program(35, 8, 0, layout) + planted = [ + _module_paths(node) + for node in _nodes_with_target(program.graph_module.graph, custom) + ] + self.mtp._rewrite_mtp_sdpa(program) + + graph = program.graph_module.graph + rewritten = _nodes_with_target( + graph, exir_ops.edge.et_vk.gemma4_sdpa.default + ) + remaining = _nodes_with_target(graph, custom) + self.assertEqual([_module_paths(node) for node in rewritten], planted) + self.assertEqual(remaining, []) + + def test_sdpa_scope_counts_fail_closed(self) -> None: + for target, assistant in ((34, 8), (36, 8), (35, 7), (35, 9)): + with self.subTest(target=target, assistant=assistant): + program = _sdpa_program(target, assistant, 0) + with self.assertRaisesRegex(ValueError, "SDPA scope mismatch"): + self.mtp._rewrite_mtp_sdpa(program) + self.assertEqual( + _nodes_with_target( + program.graph_module.graph, + exir_ops.edge.et_vk.gemma4_sdpa.default, + ), + [], + ) + + def test_unscoped_sdpa_is_rejected(self) -> None: + program = _sdpa_program(35, 8, 1) + with self.assertRaisesRegex(ValueError, "unscoped=1"): + self.mtp._rewrite_mtp_sdpa(program) + self.assertEqual( + _nodes_with_target( + program.graph_module.graph, + exir_ops.edge.et_vk.gemma4_sdpa.default, + ), + [], + ) + + def test_sdpa_abi_and_argument_guards_reject_bad_calls(self) -> None: + program = _sdpa_program(35, 8, 0) + first = _nodes_with_target( + program.graph_module.graph, exir_ops.edge.llama.custom_sdpa.default + )[0] + first.args = first.args[:7] + with self.assertRaisesRegex(ValueError, "positional ABI"): + self.mtp._rewrite_mtp_sdpa(program) + + program = _sdpa_program(35, 8, 0) + first = _nodes_with_target( + program.graph_module.graph, exir_ops.edge.llama.custom_sdpa.default + )[0] + first.args = (*first.args[:7], 0.125) + with self.assertRaisesRegex(ValueError, "not WebGPU-compatible"): + self.mtp._rewrite_mtp_sdpa(program) + + +class MTPScopedRoPETest(_MtpPartitionerTest): + def test_the_official_rope_split_is_accepted_without_rewriting(self) -> None: + rope = exir_ops.edge.et_vk.apply_rotary_emb_hf_single.default + for layout in _SCOPE_LAYOUTS: + with self.subTest(layout=layout): + program = _rope_program(20, 8, 0, layout) + planted = [ + _module_paths(node) + for node in _nodes_with_target(program.graph_module.graph, rope) + ] + + self.mtp._replace_mtp_single_hf_rope(program) + + self.assertEqual( + [ + _module_paths(node) + for node in _nodes_with_target(program.graph_module.graph, rope) + ], + planted, + ) + + def test_single_hf_rope_scope_counts_fail_closed(self) -> None: + for target, assistant, unscoped in ( + (19, 8, 0), + (21, 8, 0), + (20, 7, 0), + (20, 9, 0), + (20, 8, 1), + ): + with self.subTest(target=target, assistant=assistant, unscoped=unscoped): + program = _rope_program(target, assistant, unscoped) + with self.assertRaisesRegex( + ValueError, "single-HF-RoPE scope mismatch" + ): + self.mtp._replace_mtp_single_hf_rope(program) + + +def _topk_node( + graph: torch.fx.Graph, + *, + source_shape: tuple[int, ...] = (1, 1, 2048), + source_dtype: torch.dtype = torch.float32, + values_shape: tuple[int, ...] = (1, 1, 32), + indices_shape: tuple[int, ...] = (1, 1, 32), + indices_dtype: torch.dtype = torch.int64, + k: int = 32, + dim: int = -1, + largest: bool = True, + sorted_values: bool = True, + outputs: object | None = None, +) -> torch.fx.Node: + source = _placeholder(graph, "scores", source_shape, source_dtype) + node = graph.call_function( + exir_ops.edge.aten.topk.default, (source, k, dim, largest, sorted_values) + ) + node.meta["val"] = ( + outputs + if outputs is not None + else (_meta(values_shape), _meta(indices_shape, indices_dtype)) + ) + return node + + +def _scatter_node( + graph: torch.fx.Graph, + *, + base_shape: tuple[int, ...] = (1, 1, 262144), + base_dtype: torch.dtype = torch.float32, + dim: int = -1, + index_shape: tuple[int, ...] = (1, 1, 4096), + index_dtype: torch.dtype = torch.int64, + source_shape: tuple[int, ...] = (1, 1, 4096), + source_dtype: torch.dtype = torch.float32, + result_shape: tuple[int, ...] = (1, 1, 262144), + result_dtype: torch.dtype = torch.float32, + drop_source: bool = False, +) -> torch.fx.Node: + base = _placeholder(graph, "base", base_shape, base_dtype) + index = _placeholder(graph, "index", index_shape, index_dtype) + source = _placeholder(graph, "source", source_shape, source_dtype) + args = (base, dim, index) if drop_source else (base, dim, index, source) + node = graph.call_function(exir_ops.edge.et_vk.scatter_src_unique.default, args) + node.meta["val"] = _meta(result_shape, result_dtype) + return node + + +class MTPExtraOpFeatureTest(_MtpPartitionerTest): + def test_mtp_features_do_not_mutate_the_global_registry(self) -> None: + registry_before = dict(vulkan_supported_ops) + partitioner = self.mtp.build_webgpu_partitioner("8da4w+emb4", mode="mtp") + + self.assertEqual(vulkan_supported_ops, registry_before) + self.assertEqual(set(vulkan_supported_ops), set(registry_before)) + self.assertIn( + exir_ops.edge.aten.topk.default, partitioner._inner.extra_op_features + ) + self.assertIn( + exir_ops.edge.et_vk.scatter_src_unique.default, + partitioner._inner.extra_op_features, + ) + self.assertNotIn( + exir_ops.edge.aten.topk.default, VulkanPartitioner().extra_op_features + ) + self.assertNotIn( + exir_ops.edge.et_vk.scatter_src_unique.default, + VulkanPartitioner().extra_op_features, + ) + + def test_plain_mode_exposes_neither_scatter_nor_topk(self) -> None: + plain_allowlist = set(_webgpu_allowlist()) + plain_features = _extra_op_features() + mtp_allowlist = set(self.mtp._mtp_webgpu_allowlist()) + mtp_features = self.mtp.mtp_extra_op_features() + + for op in ( + exir_ops.edge.aten.topk.default, + exir_ops.edge.et_vk.scatter_src_unique.default, + ): + self.assertNotIn(op, plain_allowlist) + self.assertNotIn(op, plain_features) + self.assertIn(op, mtp_allowlist) + self.assertIn(op, mtp_features) + self.assertTrue(plain_allowlist.issubset(mtp_allowlist)) + self.assertEqual(len(self.mtp._mtp_webgpu_allowlist()), len(mtp_allowlist)) + + def test_uncertified_scatter_is_unreachable(self) -> None: + self.assertNotIn( + exir_ops.edge.et_vk.scatter_src_unique.default, vulkan_supported_ops + ) + self.assertFalse(has_impl(exir_ops.edge.et_vk.scatter_src_unique.default)) + self.assertNotIn(exir_ops.edge.aten.scatter.src, vulkan_supported_ops) + self.assertNotIn( + exir_ops.edge.aten.scatter.src, self.mtp.mtp_extra_op_features() + ) + self.assertNotIn( + exir_ops.edge.aten.scatter.src, set(self.mtp._mtp_webgpu_allowlist()) + ) + self.assertNotIn(exir_ops.edge.aten.topk.default, vulkan_supported_ops) + + def test_topk_gate_accepts_only_the_official_qat_shape(self) -> None: + gate = self.mtp.mtp_extra_op_features()[ + exir_ops.edge.aten.topk.default + ].are_node_inputs_supported_fn + graph = torch.fx.Graph() + self.assertTrue(gate(_topk_node(graph))) + + for case, override in { + "narrow_input": {"source_shape": (1, 1, 2047)}, + "wide_input": {"source_shape": (1, 1, 2049)}, + "rank_two_input": {"source_shape": (1, 2048)}, + "half_input": {"source_dtype": torch.float16}, + "short_values": {"values_shape": (1, 1, 31)}, + "long_values": {"values_shape": (1, 1, 33)}, + "short_indices": {"indices_shape": (1, 1, 31)}, + "float_indices": {"indices_dtype": torch.float32}, + "small_k": {"k": 31}, + "large_k": {"k": 33}, + "leading_dim": {"dim": 0}, + "non_negative_dim": {"dim": 2}, + "smallest": {"largest": False}, + "unsorted": {"sorted_values": False}, + }.items(): + with self.subTest(case=case): + self.assertFalse(gate(_topk_node(graph, **override))) + with self.subTest(case="single_output"): + self.assertFalse(gate(_topk_node(graph, outputs=_meta((1, 1, 32))))) + + def test_scatter_gate_accepts_only_the_official_qat_shape(self) -> None: + gate = self.mtp.mtp_extra_op_features()[ + exir_ops.edge.et_vk.scatter_src_unique.default + ].are_node_inputs_supported_fn + graph = torch.fx.Graph() + self.assertTrue(gate(_scatter_node(graph))) + + for case, override in { + "narrow_base": {"base_shape": (1, 1, 262143)}, + "wide_base": {"base_shape": (1, 1, 262145)}, + "half_base": {"base_dtype": torch.float16}, + "leading_dim": {"dim": 0}, + "non_negative_dim": {"dim": 2}, + "short_index": {"index_shape": (1, 1, 4095)}, + "long_index": {"index_shape": (1, 1, 4097)}, + "float_index": {"index_dtype": torch.float32}, + "short_source": {"source_shape": (1, 1, 4095)}, + "integer_source": {"source_dtype": torch.int64}, + "narrow_result": {"result_shape": (1, 1, 262143)}, + "integer_result": {"result_dtype": torch.int64}, + "missing_source": {"drop_source": True}, + }.items(): + with self.subTest(case=case): + self.assertFalse(gate(_scatter_node(graph, **override))) + + def test_mtp_partitioner_rejects_unofficial_configuration(self) -> None: + with self.assertRaisesRegex(ValueError, "8da4w\\+emb4"): + self.mtp.build_webgpu_partitioner("8da4w+emb8", mode="mtp") + with self.assertRaisesRegex(ValueError, "cannot override"): + self.mtp.build_webgpu_partitioner( + "8da4w+emb4", + mode="mtp", + compile_options={"skip_bool_tensors": True}, + ) + with self.assertRaisesRegex(ValueError, "cannot override"): + self.mtp.build_webgpu_partitioner( + "8da4w+emb4", + mode="mtp", + compile_options={"require_dynamic_shapes": False}, + ) + + +def _certified_round( + hidden: torch.Tensor, + centroid_weight: torch.Tensor, + embedding_weight: torch.Tensor, + embedding_scales: torch.Tensor, + ordering: torch.Tensor, + output_template: torch.Tensor, +) -> torch.Tensor: + scores = torch.nn.functional.linear(hidden, centroid_weight) + selected = torch.topk(scores, 32, dim=-1, largest=True, sorted=True)[1] + rows = torch.nn.functional.embedding(selected, ordering) + converted = rows.to(torch.int64) + index = converted.view(1, 1, 4096) + flat_index = converted.view(4096) + selected_embeddings = torch.ops.quantized_decomposed.embedding_4bit.dtype( + embedding_weight, + embedding_scales, + None, + -8, + 7, + flat_index, + dtype=torch.float32, + ) + selected_transpose = selected_embeddings.view(1, 1, 4096, 256).transpose(2, 3) + source = torch.matmul(hidden.unsqueeze(2), selected_transpose).squeeze(2) + return output_template.scatter(-1, index, source) + + +class _CertifiedResidualFixture(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.register_buffer("centroid_weight", torch.zeros(2048, 256)) + self.register_buffer( + "embedding_weight", torch.zeros(262144, 128, dtype=torch.int8) + ) + self.register_buffer("embedding_scales", torch.ones(262144, 1)) + self.register_buffer( + "output_template", + torch.full((1, 1, 262144), torch.finfo(torch.float32).min), + ) + + def _round(self, hidden: torch.Tensor, ordering: torch.Tensor) -> torch.Tensor: + return _certified_round( + hidden, + self.centroid_weight, + self.embedding_weight, + self.embedding_scales, + ordering, + self.output_template, + ) + + +class _CertifiedResidualChain(_CertifiedResidualFixture): + def __init__(self, ordering: torch.Tensor, persistent: bool) -> None: + super().__init__() + self.register_buffer( + "token_ordering", ordering.to(torch.float32), persistent=persistent + ) + + def forward( + self, + first_hidden: torch.Tensor, + second_hidden: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + return ( + self._round(first_hidden, self.token_ordering), + self._round(second_hidden, self.token_ordering), + ) + + +class _NonconstantOrderingChain(_CertifiedResidualFixture): + def forward( + self, + first_hidden: torch.Tensor, + second_hidden: torch.Tensor, + ordering: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + return ( + self._round(first_hidden, ordering), + self._round(second_hidden, ordering), + ) + + +def _identity_ordering() -> torch.Tensor: + return torch.arange(262144, dtype=torch.int64).reshape(2048, 128) + + +def _chain_inputs() -> tuple[torch.Tensor, ...]: + return ( + torch.zeros(1, 1, 256), + torch.ones(1, 1, 256), + ) + + +def _export_certified_chain( + persistent: bool = False, +) -> torch.export.ExportedProgram: + return torch.export.export( + _CertifiedResidualChain(_identity_ordering(), persistent), + _chain_inputs(), + strict=True, + ).run_decompositions({}) + + +def _export_nonconstant_ordering_chain() -> torch.export.ExportedProgram: + return torch.export.export( + _NonconstantOrderingChain(), + (*_chain_inputs(), _identity_ordering().to(torch.float32)), + strict=True, + ).run_decompositions({}) + + +def _mtp_delegation_tags( + program: torch.export.ExportedProgram, + partitioner: Gemma4WebGPUPartitioner, +) -> dict[object, list[object]]: + # The production-shaped matmul decomposes through non-resizable + # expand_copy nodes. This certifier unit owns op eligibility; real-model + # serializer closure is covered by the export and artifact gates. + edge = to_edge( + program, + compile_config=EdgeCompileConfig(_check_ir_validity=False), + ).exported_program() + result = partitioner.partition(edge) + targets = ( + exir_ops.edge.aten.topk.default, + exir_ops.edge.et_vk.scatter_src_unique.default, + ) + return { + target: [ + node.meta.get("delegation_tag") + for node in result.tagged_exported_program.graph.nodes + if node.target == target + ] + for target in targets + } + + +class MTPCertifiedScatterTest(_MtpPartitionerTest): + def test_exactly_two_certified_sites_are_rewritten(self) -> None: + for persistent in (False, True): + with self.subTest(persistent_ordering_buffer=persistent): + program = _export_certified_chain(persistent) + rewrites = self.mtp.rewrite_certified_unique_scatter( + program, _identity_ordering(), expected_chains=2 + ) + + self.assertEqual(rewrites, 2) + self.assertEqual( + len( + _nodes_with_target( + program.graph, + torch.ops.et_vk.scatter_src_unique.default, + ) + ), + 2, + ) + self.assertEqual( + _nodes_with_target(program.graph, torch.ops.aten.scatter.src), + [], + ) + + def test_a_rejected_rewrite_leaves_the_graph_untouched(self) -> None: + duplicated = _identity_ordering().reshape(-1).clone() + duplicated[0] = duplicated[1] + rejections = { + "chain_count_one": (_identity_ordering(), 1, "residual topology mismatch"), + "chain_count_three": ( + _identity_ordering(), + 3, + "residual topology mismatch", + ), + "chain_count_zero": ( + _identity_ordering(), + 0, + "residual topology mismatch", + ), + "foreign_ordering": ( + _identity_ordering().reshape(-1).roll(1).reshape(2048, 128), + 2, + "token-ordering identity mismatch", + ), + "non_permutation_ordering": ( + duplicated.reshape(2048, 128), + 2, + "exact permutation", + ), + } + for case, (ordering, expected_chains, message) in rejections.items(): + with self.subTest(case=case): + program = _export_certified_chain() + with self.assertRaisesRegex(ValueError, message): + self.mtp.rewrite_certified_unique_scatter( + program, ordering, expected_chains=expected_chains + ) + self.assertEqual( + len(_nodes_with_target(program.graph, torch.ops.aten.scatter.src)), + 2, + ) + + def test_a_nonconstant_ordering_source_is_rejected(self) -> None: + program = _export_nonconstant_ordering_chain() + with self.assertRaisesRegex(ValueError, "token-ordering identity mismatch"): + self.mtp.rewrite_certified_unique_scatter( + program, _identity_ordering(), expected_chains=2 + ) + self.assertEqual( + len(_nodes_with_target(program.graph, torch.ops.aten.scatter.src)), 2 + ) + + def test_mtp_partitioner_tags_certified_topk_and_scatter(self) -> None: + program = _export_certified_chain() + self.mtp.rewrite_certified_unique_scatter( + program, _identity_ordering(), expected_chains=2 + ) + tags = _mtp_delegation_tags( + program, + self.mtp.build_webgpu_partitioner("8da4w+emb4", mode="mtp"), + ) + + for op in ( + exir_ops.edge.aten.topk.default, + exir_ops.edge.et_vk.scatter_src_unique.default, + ): + self.assertEqual(len(tags[op]), 2) + self.assertTrue(all(tag is not None for tag in tags[op])) + + def test_dropping_an_mtp_feature_breaks_the_lowering_gate(self) -> None: + for op in ( + exir_ops.edge.aten.topk.default, + exir_ops.edge.et_vk.scatter_src_unique.default, + ): + with self.subTest(op=op.__name__): + program = _export_certified_chain() + self.mtp.rewrite_certified_unique_scatter( + program, _identity_ordering(), expected_chains=2 + ) + partitioner = self.mtp.build_webgpu_partitioner( + "8da4w+emb4", mode="mtp" + ) + del partitioner._inner.extra_op_features[op] + tags = _mtp_delegation_tags(program, partitioner) + self.assertEqual(tags[op], [None, None]) + other = ( + exir_ops.edge.et_vk.scatter_src_unique.default + if op == exir_ops.edge.aten.topk.default + else exir_ops.edge.aten.topk.default + ) + self.assertEqual(len(tags[other]), 2) + self.assertTrue(all(tag is not None for tag in tags[other])) diff --git a/examples/models/gemma4/tests/test_gemma4_spec_runner_contract.cpp b/examples/models/gemma4/tests/test_gemma4_spec_runner_contract.cpp new file mode 100644 index 00000000000..67a453ee579 --- /dev/null +++ b/examples/models/gemma4/tests/test_gemma4_spec_runner_contract.cpp @@ -0,0 +1,408 @@ +/* + * 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 + +namespace executorch::examples::gemma4 { +namespace { + +constexpr int64_t kVocabSize = 262144; +const std::vector kStopTokens = {1, 106, 50}; +const std::vector kNoStopTokens = {}; + +// Two accepted drafts plus the bonus: the widest well-formed K=2 round. +Gemma4K2Output FullMatch() { + return Gemma4K2Output{{10, 11}, {10, 11, 90}, 2, 90, 0.0f}; +} + +bool Rejected(const Gemma4K2Output& output) { + return !reconcile_gemma4_k2(output, 2, 3, kStopTokens).valid; +} + +TEST(Gemma4SpecControllerTest, AdvancesByAcceptedPrefixAndSeedsFromBonus) { + const auto no_match = reconcile_gemma4_k2( + Gemma4K2Output{{10, 11}, {90, 91, 92}, 0, 90, 0.0f}, 2, 3, kStopTokens); + EXPECT_TRUE(no_match.valid); + EXPECT_EQ(no_match.committed, std::vector({90})); + EXPECT_EQ(no_match.selected, std::vector({90})); + EXPECT_EQ(no_match.discarded, std::vector({})); + EXPECT_EQ(no_match.accepted_drafts, 0u); + EXPECT_EQ(no_match.next_position, 3); + EXPECT_EQ(no_match.next_seed, 90); + EXPECT_FALSE(no_match.stopped); + + const auto one_match = reconcile_gemma4_k2( + Gemma4K2Output{{10, 11}, {10, 90, 92}, 1, 90, 0.0f}, 2, 3, kStopTokens); + EXPECT_TRUE(one_match.valid); + EXPECT_EQ(one_match.committed, std::vector({10, 90})); + EXPECT_EQ(one_match.selected, std::vector({10, 90})); + EXPECT_EQ(one_match.accepted_drafts, 1u); + EXPECT_EQ(one_match.next_position, 4); + EXPECT_EQ(one_match.next_seed, 90); + + const auto two_matches = reconcile_gemma4_k2(FullMatch(), 2, 3, kStopTokens); + EXPECT_TRUE(two_matches.valid); + EXPECT_EQ(two_matches.committed, std::vector({10, 11, 90})); + EXPECT_EQ(two_matches.selected, std::vector({10, 11, 90})); + EXPECT_EQ(two_matches.accepted_drafts, 2u); + EXPECT_EQ(two_matches.next_position, 5); + EXPECT_EQ(two_matches.next_seed, 90); +} + +TEST(Gemma4SpecControllerTest, ChainedRoundsWalkStartPositionsTwoThreeFive) { + const auto first = reconcile_gemma4_k2( + Gemma4K2Output{{10, 11}, {90, 91, 92}, 0, 90, 0.0f}, 2, 8, kStopTokens); + ASSERT_TRUE(first.valid); + EXPECT_EQ(first.next_position, 3); + EXPECT_EQ(first.next_seed, 90); + + const auto second = reconcile_gemma4_k2( + Gemma4K2Output{{20, 21}, {20, 91, 92}, 1, 91, 0.0f}, + first.next_position, + 8, + kStopTokens); + ASSERT_TRUE(second.valid); + EXPECT_EQ(second.next_position, 5); + EXPECT_EQ(second.next_seed, 91); + + const auto third = reconcile_gemma4_k2( + Gemma4K2Output{{30, 31}, {30, 31, 92}, 2, 92, 0.0f}, + second.next_position, + 8, + kStopTokens); + ASSERT_TRUE(third.valid); + EXPECT_EQ(third.next_position, 8); + EXPECT_EQ(third.next_seed, 92); +} + +TEST(Gemma4SpecControllerTest, SeedsFromBonusNotFromDraftOrTargetTail) { + const auto decision = reconcile_gemma4_k2( + Gemma4K2Output{{10, 11}, {10, 90, 92}, 1, 90, 0.0f}, 2, 3, kStopTokens); + ASSERT_TRUE(decision.valid); + EXPECT_EQ(decision.next_seed, 90); + EXPECT_NE(decision.next_seed, 10); + EXPECT_NE(decision.next_seed, 11); + EXPECT_NE(decision.next_seed, 92); + EXPECT_EQ(decision.selected.back(), decision.next_seed); +} + +TEST(Gemma4SpecControllerTest, TruncatesAtStopWithoutCommittingTheStopToken) { + const auto decision = reconcile_gemma4_k2( + Gemma4K2Output{{106, 11}, {106, 11, 90}, 2, 90, 0.0f}, 2, 3, kStopTokens); + + EXPECT_TRUE(decision.valid); + EXPECT_TRUE(decision.stopped); + EXPECT_EQ(decision.stop_token, 106); + EXPECT_TRUE(decision.committed.empty()); + EXPECT_EQ(decision.discarded, std::vector({11, 90})); + EXPECT_EQ(decision.next_position, 5); + EXPECT_EQ(decision.next_seed, 90); +} + +TEST(Gemma4SpecControllerTest, StopTokenInBonusSlotCommitsAcceptedPrefix) { + const auto decision = reconcile_gemma4_k2( + Gemma4K2Output{{10, 11}, {10, 11, 1}, 2, 1, 0.0f}, 2, 3, kStopTokens); + + EXPECT_TRUE(decision.valid); + EXPECT_TRUE(decision.stopped); + EXPECT_EQ(decision.stop_token, 1); + EXPECT_EQ(decision.committed, std::vector({10, 11})); + EXPECT_TRUE(decision.discarded.empty()); +} + +TEST(Gemma4SpecControllerTest, TruncatesAtBudgetWithoutAnotherRound) { + const auto one = reconcile_gemma4_k2(FullMatch(), 2, 1, kStopTokens); + EXPECT_TRUE(one.valid); + EXPECT_FALSE(one.stopped); + EXPECT_EQ(one.committed, std::vector({10})); + EXPECT_EQ(one.discarded, std::vector({11, 90})); + + const auto two = reconcile_gemma4_k2(FullMatch(), 2, 2, kStopTokens); + EXPECT_EQ(two.committed, std::vector({10, 11})); + EXPECT_EQ(two.discarded, std::vector({90})); + + const auto three = reconcile_gemma4_k2(FullMatch(), 2, 3, kStopTokens); + EXPECT_EQ(three.committed, std::vector({10, 11, 90})); + EXPECT_TRUE(three.discarded.empty()); +} + +TEST(Gemma4SpecControllerTest, RejectsInconsistentGraphMatchMetadata) { + EXPECT_TRUE(Rejected(Gemma4K2Output{{10, 11}, {10, 91, 92}, 2, 92, 0.0f})); + EXPECT_TRUE(Rejected(Gemma4K2Output{{10, 11}, {90, 91, 92}, 1, 91, 0.0f})); + EXPECT_TRUE(Rejected(Gemma4K2Output{{10, 11}, {10, 11, 90}, 1, 11, 0.0f})); + EXPECT_TRUE(Rejected(Gemma4K2Output{{10, 11}, {10, 90, 92}, 0, 10, 0.0f})); +} + +TEST(Gemma4SpecControllerTest, RejectsBonusThatIsNotGreedyAtMatchCount) { + EXPECT_TRUE(Rejected(Gemma4K2Output{{10, 11}, {10, 11, 90}, 2, 11, 0.0f})); + EXPECT_TRUE(Rejected(Gemma4K2Output{{10, 11}, {90, 91, 92}, 0, 91, 0.0f})); + EXPECT_TRUE(Rejected(Gemma4K2Output{{10, 11}, {10, 90, 92}, 1, 92, 0.0f})); +} + +TEST(Gemma4SpecControllerTest, RejectsStartPositionBelowTwo) { + for (const int64_t start : {-1, 0, 1}) { + EXPECT_FALSE(reconcile_gemma4_k2(FullMatch(), start, 3, kStopTokens).valid) + << "start_position=" << start; + } + EXPECT_TRUE(reconcile_gemma4_k2(FullMatch(), 2, 3, kStopTokens).valid); +} + +TEST(Gemma4SpecControllerTest, RejectsZeroTokenBudget) { + EXPECT_FALSE(reconcile_gemma4_k2(FullMatch(), 2, 0, kStopTokens).valid); + EXPECT_TRUE(reconcile_gemma4_k2(FullMatch(), 2, 1, kStopTokens).valid); +} + +TEST(Gemma4SpecControllerTest, RejectsNonPositiveVocabSize) { + EXPECT_FALSE(reconcile_gemma4_k2(FullMatch(), 2, 3, kStopTokens, 0).valid); + EXPECT_FALSE(reconcile_gemma4_k2(FullMatch(), 2, 3, kStopTokens, -1).valid); + EXPECT_TRUE(reconcile_gemma4_k2(FullMatch(), 2, 3, kStopTokens, 91).valid); +} + +TEST(Gemma4SpecControllerTest, RejectsMatchCountOutsideZeroToTwo) { + Gemma4K2Output low = FullMatch(); + low.match_count = -1; + EXPECT_FALSE(reconcile_gemma4_k2(low, 2, 3, kStopTokens).valid); + + Gemma4K2Output high = FullMatch(); + high.match_count = 3; + EXPECT_FALSE(reconcile_gemma4_k2(high, 2, 3, kStopTokens).valid); +} + +TEST(Gemma4SpecControllerTest, RejectsNonFiniteStateProbe) { + for (const float probe : + {std::numeric_limits::quiet_NaN(), + std::numeric_limits::infinity(), + -std::numeric_limits::infinity()}) { + Gemma4K2Output output = FullMatch(); + output.state_probe = probe; + EXPECT_FALSE(reconcile_gemma4_k2(output, 2, 3, kStopTokens).valid); + } + Gemma4K2Output finite = FullMatch(); + finite.state_probe = -3.5f; + EXPECT_TRUE(reconcile_gemma4_k2(finite, 2, 3, kStopTokens).valid); +} + +TEST(Gemma4SpecControllerTest, RejectsOutOfRangeTokenIds) { + Gemma4K2Output low_candidate = FullMatch(); + low_candidate.candidates = {-1, 11}; + EXPECT_FALSE(reconcile_gemma4_k2(low_candidate, 2, 3, kStopTokens).valid); + + Gemma4K2Output high_candidate = FullMatch(); + high_candidate.candidates = {kVocabSize, 11}; + EXPECT_FALSE(reconcile_gemma4_k2(high_candidate, 2, 3, kStopTokens).valid); + + Gemma4K2Output low_greedy = FullMatch(); + low_greedy.target_greedy = {10, 11, -1}; + EXPECT_FALSE(reconcile_gemma4_k2(low_greedy, 2, 3, kStopTokens).valid); + + Gemma4K2Output high_greedy = FullMatch(); + high_greedy.target_greedy = {10, 11, kVocabSize}; + high_greedy.bonus = kVocabSize; + EXPECT_FALSE(reconcile_gemma4_k2(high_greedy, 2, 3, kStopTokens).valid); + + Gemma4K2Output narrow_vocab = FullMatch(); + EXPECT_FALSE(reconcile_gemma4_k2(narrow_vocab, 2, 3, kStopTokens, 90).valid); +} + +TEST(Gemma4SpecControllerTest, EmptyStopTokenListNeverStops) { + const auto decision = reconcile_gemma4_k2( + Gemma4K2Output{{106, 1}, {106, 1, 90}, 2, 90, 0.0f}, 2, 3, kNoStopTokens); + EXPECT_TRUE(decision.valid); + EXPECT_FALSE(decision.stopped); + EXPECT_EQ(decision.stop_token, -1); + EXPECT_EQ(decision.committed, std::vector({106, 1, 90})); +} + +TEST(Gemma4SpecControllerTest, RejectedDecisionKeepsDocumentedDefaults) { + const auto decision = reconcile_gemma4_k2( + Gemma4K2Output{{10, 11}, {10, 91, 92}, 2, 92, 0.0f}, 2, 3, kStopTokens); + EXPECT_FALSE(decision.valid); + EXPECT_FALSE(decision.stopped); + EXPECT_EQ(decision.stop_token, -1); + EXPECT_EQ(decision.next_position, -1); + EXPECT_EQ(decision.next_seed, -1); + EXPECT_EQ(decision.accepted_drafts, 0u); + EXPECT_TRUE(decision.selected.empty()); + EXPECT_TRUE(decision.committed.empty()); + EXPECT_TRUE(decision.discarded.empty()); +} + +TEST(Gemma4SpecControllerTest, DefaultVocabSizeMatchesTheExportContract) { + Gemma4K2Output output = FullMatch(); + output.candidates = {kVocabSize - 1, 11}; + output.target_greedy = {kVocabSize - 1, 11, 90}; + EXPECT_TRUE(reconcile_gemma4_k2(output, 2, 3, kStopTokens).valid); + + output.candidates = {kVocabSize, 11}; + output.target_greedy = {kVocabSize, 11, 90}; + EXPECT_FALSE(reconcile_gemma4_k2(output, 2, 3, kStopTokens).valid); +} + +TEST(Gemma4SpecControllerTest, ConfigDefaultsMatchTheExportContract) { + const Gemma4SpecRunnerConfig config; + EXPECT_EQ(config.vocab_size, kVocabSize); + EXPECT_EQ(config.max_input_length, 512); + EXPECT_EQ(config.target_capacity, 8960); + EXPECT_EQ(config.donor_capacity, 8960); + EXPECT_EQ(config.method_name, "k2_round"); +} + +using ::executorch::runtime::Error; + +TEST(Gemma4SpecRunnerLifecycleTest, FreshRunnerReportsEmptyAccounting) { + Gemma4SpecRunner runner; + EXPECT_FALSE(runner.is_loaded()); + EXPECT_EQ(runner.execute_count(), 0u); + EXPECT_EQ(runner.accepted_drafts(), 0u); + EXPECT_EQ(runner.buffered_tokens(), 0u); +} + +TEST(Gemma4SpecRunnerLifecycleTest, ResetWithoutAModuleIsInvalidState) { + Gemma4SpecRunner runner; + EXPECT_EQ(runner.reset(), Error::InvalidState); + EXPECT_EQ(runner.reset(), Error::InvalidState); + EXPECT_FALSE(runner.is_loaded()); + EXPECT_EQ(runner.buffered_tokens(), 0u); + EXPECT_EQ(runner.execute_count(), 0u); + EXPECT_EQ(runner.accepted_drafts(), 0u); +} + +TEST(Gemma4SpecRunnerLifecycleTest, UnloadWithoutAModuleSucceeds) { + Gemma4SpecRunner runner; + EXPECT_EQ(runner.unload(), Error::Ok); + EXPECT_FALSE(runner.is_loaded()); +} + +TEST(Gemma4SpecRunnerLifecycleTest, LoadRejectsEmptyPathAndInvalidConfig) { + Gemma4SpecRunner runner; + EXPECT_EQ(runner.load("", {}), Error::InvalidArgument); + + Gemma4SpecRunnerConfig config; + config.vocab_size = 0; + Gemma4SpecRunner zero_vocab(config); + EXPECT_EQ(zero_vocab.load("model.pte", {}), Error::InvalidArgument); + + Gemma4SpecRunnerConfig unnamed; + unnamed.method_name = ""; + Gemma4SpecRunner no_method(unnamed); + EXPECT_EQ(no_method.load("model.pte", {}), Error::InvalidArgument); +} + +// Earlier empty-PTD cases cannot isolate a config clause. These pass exactly +// three PTDs and vary one clause at a time. +TEST(Gemma4SpecRunnerLifecycleTest, LoadRejectsEachInvalidConfigClause) { + const auto rejects = [](const char* clause, + const Gemma4SpecRunnerConfig& config) { + Gemma4SpecRunner runner(config); + EXPECT_EQ( + runner.load("k2_round.pte", {"a.ptd", "b.ptd", "c.ptd"}), + Error::InvalidArgument) + << clause; + EXPECT_FALSE(runner.is_loaded()) << clause; + }; + const int64_t above_int32 = + static_cast(std::numeric_limits::max()) + 1; + + for (const int64_t vocab : {int64_t{0}, int64_t{-1}}) { + Gemma4SpecRunnerConfig config; + config.vocab_size = vocab; + rejects("vocab_size", config); + } + for (const int64_t length : {int64_t{0}, int64_t{-1}, above_int32}) { + Gemma4SpecRunnerConfig config; + config.max_input_length = length; + rejects("max_input_length", config); + } + for (const int64_t capacity : {int64_t{0}, int64_t{-1}}) { + Gemma4SpecRunnerConfig target; + target.target_capacity = capacity; + rejects("target_capacity", target); + Gemma4SpecRunnerConfig donor; + donor.donor_capacity = capacity; + rejects("donor_capacity", donor); + } + for (const char* name : {"", "k2"}) { + Gemma4SpecRunnerConfig config; + config.method_name = name; + rejects("method_name", config); + } +} + +TEST(Gemma4SpecRunnerLifecycleTest, LoadRejectsAnyPtdCountOtherThanThree) { + const std::vector> wrong = { + {}, {"a.ptd"}, {"a.ptd", "b.ptd"}, {"a.ptd", "b.ptd", "c.ptd", "d.ptd"}}; + for (const auto& ptds : wrong) { + Gemma4SpecRunner runner; + EXPECT_EQ(runner.load("k2_round.pte", ptds), Error::InvalidArgument) + << ptds.size() << " PTDs"; + } +} + +TEST(Gemma4SpecRunnerLifecycleTest, WellFormedLoadIsNotRejectedByArgumentFence) { + Gemma4SpecRunner runner; + const Error error = + runner.load("no_such_k2_round.pte", {"a.ptd", "b.ptd", "c.ptd"}); + + EXPECT_NE(error, Error::InvalidArgument); + EXPECT_NE(error, Error::Ok); + EXPECT_FALSE(runner.is_loaded()); +} +TEST(Gemma4SpecRunnerLifecycleTest, UnloadedRunnerRejectsEveryStepEntry) { + Gemma4SpecRunner runner; + EXPECT_EQ( + runner.execute({10}, {2}, false, 2).error(), Error::InvalidState); + EXPECT_EQ(runner.prefill_step(10, 0), Error::InvalidArgument); + EXPECT_EQ(runner.step(10, 2).error(), Error::InvalidArgument); +} + +TEST(Gemma4SpecRunnerLifecycleTest, RejectedExecutionsBillNothingAcrossReset) { + Gemma4SpecRunner runner; + const auto first = runner.execute({10, 11}, {2, 3}, false, 2); + EXPECT_EQ(first.error(), Error::InvalidState); + EXPECT_EQ(runner.execute_count(), 0u); + + EXPECT_EQ(runner.reset(), Error::InvalidState); + + const auto second = runner.execute({10, 11}, {2, 3}, false, 2); + EXPECT_EQ(second.error(), first.error()); + EXPECT_EQ(runner.execute_count(), 0u); + EXPECT_EQ(runner.accepted_drafts(), 0u); + EXPECT_EQ(runner.buffered_tokens(), 0u); + EXPECT_FALSE(runner.is_loaded()); +} + +TEST(Gemma4SpecRunnerLifecycleTest, GenerateRejectsMalformedRequests) { + Gemma4SpecRunner runner; + EXPECT_EQ(runner.generate({}, 4, {}).error(), Error::InvalidArgument); + EXPECT_EQ(runner.generate({10, 11}, 0, {}).error(), Error::InvalidArgument); + EXPECT_EQ(runner.generate({10}, 2, {}).error(), Error::InvalidArgument); + EXPECT_EQ( + runner.generate({10, 11}, 4, {kVocabSize}).error(), + Error::InvalidArgument); + EXPECT_EQ(runner.generate({10, 11}, 4, {-1}).error(), Error::InvalidArgument); + EXPECT_EQ(runner.generate({10, 11}, 4, {}).error(), Error::InvalidState); +} + +TEST(Gemma4SpecRunnerLifecycleTest, ProfileJsonIsSchemaVersionOne) { + Gemma4SpecRunner runner; + runner.set_profiling_enabled(true); + const std::string profile = runner.profile_json(); + EXPECT_NE(profile.find("\"schemaVersion\":1"), std::string::npos); + EXPECT_NE(profile.find("\"execute_generation\":0"), std::string::npos); + runner.set_profiling_enabled(false); + EXPECT_EQ(runner.profile_json(), profile); +} + +} // namespace +} // namespace executorch::examples::gemma4 diff --git a/examples/models/gemma4/tests/test_mtp_spec_oracle.py b/examples/models/gemma4/tests/test_mtp_spec_oracle.py new file mode 100644 index 00000000000..df6a23dc973 --- /dev/null +++ b/examples/models/gemma4/tests/test_mtp_spec_oracle.py @@ -0,0 +1,1210 @@ +# 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. + +# pyre-strict + +from __future__ import annotations + +import ast +import copy +import hashlib +import importlib.util +import inspect +import json +import os +import re +import tempfile +import unittest + +from pathlib import Path +from typing import Any, Mapping +from unittest import mock + +import torch + +from executorch.examples.models.gemma4.export_assistant_webgpu_artifacts import ( + _tensor_sha256 as assistant_tensor_sha256, +) +from executorch.examples.models.gemma4.generate_target_prefill_oracle import ( + _tensor_bytes as target_prefill_tensor_bytes, +) +from executorch.examples.models.gemma4.target_prefill_contract import ( + canonical_json_bytes, + final_chunk_range, + prompt_plan_sha256, + reviewed_producer_source_path, + TARGET_PREFILL_AUTHORITY, + TARGET_PREFILL_CONTEXTS, + TARGET_PREFILL_ENVELOPE_KIND, + TARGET_PREFILL_SCHEMA_VERSION, +) +from executorch.examples.models.gemma4.tests import ( + generate_mtp_spec_oracle as mtp_oracle, +) +from executorch.examples.models.gemma4.tests.generate_mtp_spec_oracle import ( + _build_parser, + assemble_oracle_document, + build_oracle_binding, + K2_ROUND_ABI, + K2_VOCAB_SIZE, + ORACLE_CONTEXT_KEYS, + ORACLE_TOKEN_BUDGET, + OracleError, + parse_contexts, + production_oracle_is_acceptable, + reconcile_k2_round, + require_prefill_token_match, + SUPPORTED_AUTHORITIES, + TARGET_PREFILL_AUTHORITIES, + validate_k2_abi_edge_census, +) +from executorch.examples.models.gemma4.webgpu_artifact_manifest import ( + ASSISTANT_CHECKPOINT_ACQUISITION, + CHECKPOINT_ACQUISITION, + MTP_EDGE_CENSUS, + MTP_SOURCE_VERIFIED_PROVENANCE, +) + + +CONTRACT_CONTEXTS = "128,511,512,513,514,1024,2048,4096,4097,8192" +CONTRACT_CONTEXT_VALUES = [128, 511, 512, 513, 514, 1024, 2048, 4096, 4097, 8192] +STOP_TOKENS = [1, 106] + + +def _witness(index: int) -> dict[str, Any]: + return { + "layer0_av_sha256": f"{index:064x}", + "layer0_qk_sha256": f"{index + 1:064x}", + "logits_sha256": f"{index + 2:064x}", + "prefill_token": 1000 + index, + } + + +def _tensor_envelope(shape: list[int], digest_digit: int) -> dict[str, Any]: + return { + "byte_order": "little", + "dtype": "float32", + "layout": "row_major_contiguous", + "sha256": f"{digest_digit:064x}", + "shape": shape, + } + + +def _target_prefill_v2_context(context: int, index: int) -> dict[str, Any]: + final_start, final_length = final_chunk_range(context) + arm_config = { + "dtype": "float32", + "enable_dynamic_shape": True, + "group_size": 128, + "max_seq_len": 8960, + "text_quantize": "8da4w+emb4", + "use_kv_cache": True, + "variant": "e2b", + } + return { + "arm_configs": { + "custom_sdpa_fused": {**arm_config, "use_custom_sdpa": True}, + "manual_unfused": {**arm_config, "use_custom_sdpa": False}, + }, + "cache_reset_counts": { + "custom_sdpa_fused": 35, + "manual_unfused": 35, + }, + "chunk_size": 512, + "context": context, + "final_chunk_length": final_length, + "final_chunk_start": final_start, + "layer0_manual_unfused_vs_custom_sdpa_fused": { + "agreement": { + "atol": 1e-4, + "max_abs": 0.0, + "passed": True, + "rel_rms": 0.0, + "rtol": 1e-3, + }, + "custom_sdpa_fused": _tensor_envelope([1, final_length, 8, 256], index + 3), + "manual_unfused": _tensor_envelope([1, final_length, 8, 256], index + 4), + }, + "logits_post_softcap": _tensor_envelope([1, 1, K2_VOCAB_SIZE], index + 2), + "logits_pre_softcap": _tensor_envelope([1, 1, K2_VOCAB_SIZE], index + 1), + "prefill_token_post_softcap": 1000 + index, + "prefill_token_raw": 1000 + index, + "prompt_plan_sha256": prompt_plan_sha256(context), + } + + +class MtpSpecOracleGeneratorTest(unittest.TestCase): + def setUp(self) -> None: + directory = tempfile.TemporaryDirectory() + self.addCleanup(directory.cleanup) + self.root = Path(directory.name) + (self.root / "assistant").mkdir() + (self.root / "target").mkdir() + (self.root / "target" / "generation_config.json").write_text( + json.dumps({"eos_token_id": STOP_TOKENS}), encoding="utf-8" + ) + self.manifest_path = self.root / "mtp-manifest.json" + self.receipt_path = self.root / "target_prefill_oracle.json" + self.output = self.root / "mtp_spec_oracle.json" + self.manifest: dict[str, Any] = { + "abi": copy.deepcopy(K2_ROUND_ABI), + "artifacts": [ + {"path": "model.pte", "role": "pte"}, + {"path": "model0.ptd", "role": "ptd"}, + {"path": "model1.ptd", "role": "ptd"}, + {"path": "model2.ptd", "role": "ptd"}, + ], + "checkpoints": {"assistant": "assistant", "target": "target"}, + "max_context_length": 8960, + "method": "k2_round", + "ptd_order": ["model0.ptd", "model1.ptd", "model2.ptd"], + "schema_version": 1, + } + self.receipt: dict[str, Any] = { + "authority": "portable_eager", + "contexts": { + str(context): _witness(index) + for index, context in enumerate(CONTRACT_CONTEXT_VALUES) + }, + "schema_version": 1, + } + self._write() + + def _write(self) -> None: + self.manifest_path.write_text(json.dumps(self.manifest), encoding="utf-8") + self.receipt_path.write_text(json.dumps(self.receipt), encoding="utf-8") + + def _build(self, contexts: str = CONTRACT_CONTEXTS) -> dict[str, Any]: + return build_oracle_binding( + self.manifest_path, + self.receipt_path, + "portable_eager", + contexts, + self.output, + ) + + def _expect_rejection(self, contexts: str = CONTRACT_CONTEXTS) -> None: + self._write() + with self.assertRaises(OracleError): + self._build(contexts) + + def _records(self, binding: dict[str, Any]) -> dict[str, dict[str, Any]]: + records: dict[str, dict[str, Any]] = {} + for context in binding["contexts"]: + key = str(context) + record: dict[str, Any] = {name: [] for name in ORACLE_CONTEXT_KEYS} + record["decoded_text"] = f"context {key}" + record["target_prefill"] = binding["target_prefill"][key] + records[key] = record + return records + + def test_binding_binds_manifest_receipt_and_contract_contexts(self) -> None: + binding = self._build() + self.assertEqual(binding["contexts"], CONTRACT_CONTEXT_VALUES) + self.assertEqual(binding["authority"], "portable_eager") + self.assertEqual(binding["max_context_length"], 8960) + self.assertEqual(binding["stop_tokens"], STOP_TOKENS) + self.assertEqual(binding["token_budget"], ORACLE_TOKEN_BUDGET) + self.assertEqual(binding["checkpoints"]["target"], self.root / "target") + self.assertEqual(binding["checkpoints"]["assistant"], self.root / "assistant") + self.assertEqual( + binding["mtp_manifest_sha256"], + hashlib.sha256(self.manifest_path.read_bytes()).hexdigest(), + ) + self.assertEqual( + binding["target_prefill_oracle_sha256"], + hashlib.sha256(self.receipt_path.read_bytes()).hexdigest(), + ) + + def test_target_prefill_witnesses_are_copied_verbatim(self) -> None: + binding = self._build() + self.assertEqual( + sorted(binding["target_prefill"]), + sorted(str(context) for context in CONTRACT_CONTEXT_VALUES), + ) + for index, context in enumerate(CONTRACT_CONTEXT_VALUES): + self.assertEqual(binding["target_prefill"][str(context)], _witness(index)) + + def test_rejects_reinterpreted_or_augmented_witness(self) -> None: + self.receipt["contexts"]["512"]["extra"] = 1 + self._expect_rejection() + + def test_assembled_document_binds_every_receipt_and_abi_identity(self) -> None: + binding = self._build() + document = assemble_oracle_document(binding, self._records(binding)) + self.assertEqual(document["schema_version"], 1) + self.assertEqual(document["method"], "k2_round") + self.assertEqual(document["abi"], K2_ROUND_ABI) + self.assertEqual(document["authority"], "portable_eager") + self.assertEqual(document["contexts"], CONTRACT_CONTEXT_VALUES) + self.assertEqual(document["stop_tokens"], STOP_TOKENS) + self.assertEqual(document["token_budget"], ORACLE_TOKEN_BUDGET) + self.assertEqual( + document["mtp_manifest_sha256"], binding["mtp_manifest_sha256"] + ) + self.assertEqual( + document["target_prefill_oracle_sha256"], + binding["target_prefill_oracle_sha256"], + ) + self.assertEqual( + sorted(document["records"]), + sorted(str(context) for context in CONTRACT_CONTEXT_VALUES), + ) + self.assertEqual(tuple(sorted(document["records"]["512"])), ORACLE_CONTEXT_KEYS) + + def test_assemble_rejects_reinterpreted_target_prefill(self) -> None: + binding = self._build() + records = self._records(binding) + records["512"]["target_prefill"] = dict(records["512"]["target_prefill"]) + records["512"]["target_prefill"]["prefill_token"] += 1 + with self.assertRaises(OracleError): + assemble_oracle_document(binding, records) + + def test_assemble_rejects_missing_extra_and_malformed_records(self) -> None: + binding = self._build() + missing = self._records(binding) + del missing["512"] + with self.assertRaises(OracleError): + assemble_oracle_document(binding, missing) + extra = self._records(binding) + extra["99"] = extra["512"] + with self.assertRaises(OracleError): + assemble_oracle_document(binding, extra) + malformed = self._records(binding) + del malformed["512"]["kv_witnesses"] + with self.assertRaises(OracleError): + assemble_oracle_document(binding, malformed) + + def _expect_authority_rejection(self, authority: str) -> None: + with self.assertRaises(OracleError) as caught: + build_oracle_binding( + self.manifest_path, + self.receipt_path, + authority, + CONTRACT_CONTEXTS, + self.output, + ) + self.assertIn(f"unknown --authority {authority!r}", str(caught.exception)) + self.assertIn("portable_eager", str(caught.exception)) + + def test_rejects_unknown_authority(self) -> None: + for authority in ("", "portable", "PORTABLE_EAGER", "webgpu_live"): + with self.subTest(authority=authority): + self._expect_authority_rejection(authority) + + def test_unknown_authority_cannot_self_certify_via_a_matching_receipt(self) -> None: + self.receipt["authority"] = "webgpu_live" + self._write() + self._expect_authority_rejection("webgpu_live") + + def test_rejects_manifest_schema_version(self) -> None: + for version in (None, 0, 2, True, "1", 1.0): + with self.subTest(version=version): + if version is None: + del self.manifest["schema_version"] + else: + self.manifest["schema_version"] = version + self._expect_rejection() + self.manifest["schema_version"] = 1 + + def test_rejects_receipt_schema_version(self) -> None: + for version in (None, 2, True, "1"): + with self.subTest(version=version): + if version is None: + del self.receipt["schema_version"] + else: + self.receipt["schema_version"] = version + self._expect_rejection() + self.receipt["schema_version"] = 1 + + def test_rejects_receipt_authority_mismatch(self) -> None: + self.receipt["authority"] = "webgpu_live" + self._expect_rejection() + + def test_rejects_mismatched_abi_manifest(self) -> None: + mutations = ( + ("buffer_mutation_count", 30), + ("seed_mutation_count", 0), + ("user_inputs", ["input_pos", "input_ids", "is_round", "donor_length"]), + ) + for key, value in mutations: + with self.subTest(key=key): + self.manifest["abi"] = copy.deepcopy(K2_ROUND_ABI) + self.manifest["abi"][key] = value + self._expect_rejection() + self.manifest["abi"] = copy.deepcopy(K2_ROUND_ABI) + self.manifest["abi"]["operator_counts"]["aten.topk.default"] = 1 + self._expect_rejection() + self.manifest["abi"] = copy.deepcopy(K2_ROUND_ABI) + self.manifest["abi"]["user_outputs"][4]["dtype"] = "int64" + self._expect_rejection() + self.manifest["abi"] = copy.deepcopy(K2_ROUND_ABI) + del self.manifest["abi"]["operator_counts"] + self._expect_rejection() + + def test_rejects_wrong_method_name(self) -> None: + for method in (None, "k1_round", "text_decoder"): + with self.subTest(method=method): + self.manifest["method"] = method + self._expect_rejection() + self.manifest["method"] = "k2_round" + + def test_rejects_artifact_role_counts(self) -> None: + original = copy.deepcopy(self.manifest["artifacts"]) + self.manifest["artifacts"] = original + [{"path": "second.pte", "role": "pte"}] + self._expect_rejection() + self.manifest["artifacts"] = original[:3] + self.manifest["ptd_order"] = ["model0.ptd", "model1.ptd"] + self._expect_rejection() + self.manifest["artifacts"] = original + self.manifest["ptd_order"] = ["model0.ptd", "model1.ptd"] + self._expect_rejection() + + def test_rejects_max_context_length(self) -> None: + for value in (None, 0, -1, 8961, True, "8960"): + with self.subTest(value=value): + self.manifest["max_context_length"] = value + self._expect_rejection() + self.manifest["max_context_length"] = 8960 + + def _expect_rejection_message(self, fragment: str) -> None: + with self.assertRaises(OracleError) as caught: + self._build() + self.assertIn(fragment, str(caught.exception)) + + def test_rejects_unreadable_manifest_and_receipt(self) -> None: + manifest, receipt = self.manifest_path, self.receipt_path + manifest.unlink() + self._expect_rejection_message(f"unreadable MTP manifest: {manifest}") + manifest.write_text("{not json", encoding="utf-8") + self._expect_rejection_message(f"malformed MTP manifest: {manifest}") + manifest.write_text("[]", encoding="utf-8") + self._expect_rejection_message( + f"MTP manifest must be a JSON object: {manifest}" + ) + self._write() + receipt.unlink() + self._expect_rejection_message(f"unreadable target-prefill receipt: {receipt}") + receipt.write_text("[1, 2]", encoding="utf-8") + self._expect_rejection_message( + f"target-prefill receipt must be a JSON object: {receipt}" + ) + + def test_rejects_existing_output(self) -> None: + self.output.write_text("{}", encoding="utf-8") + with self.assertRaises(OracleError): + self._build() + self.output.unlink() + self.output.symlink_to(self.root / "missing.json") + with self.assertRaises(OracleError): + self._build() + + def test_rejects_broken_checkpoint_binding(self) -> None: + for checkpoints in ( + None, + {}, + {"target": "target"}, + {"assistant": "assistant", "target": "target", "extra": "target"}, + {"assistant": "assistant", "target": "missing"}, + {"assistant": "assistant", "target": ""}, + {"assistant": "assistant", "target": 3}, + ): + with self.subTest(checkpoints=checkpoints): + self.manifest["checkpoints"] = checkpoints + self._expect_rejection() + + def test_rejects_invalid_stop_tokens(self) -> None: + config = self.root / "target" / "generation_config.json" + for value in ([], [1, 1], [-1], [K2_VOCAB_SIZE], ["1"], [True]): + with self.subTest(value=value): + config.write_text(json.dumps({"eos_token_id": value}), encoding="utf-8") + self._expect_rejection() + config.write_text(json.dumps({}), encoding="utf-8") + self._expect_rejection() + config.unlink() + self._expect_rejection() + + def test_rejects_malformed_context_lists(self) -> None: + for contexts in ( + "", + " ", + "128,", + ",128", + "128,,512", + "-1", + "0", + "1.5", + "128,128", + "8961", + "0x80", + "128 512", + ): + with self.subTest(contexts=contexts): + with self.assertRaises(OracleError): + self._build(contexts) + + def test_parse_contexts_accepts_the_command_contract_list(self) -> None: + self.assertEqual( + parse_contexts(CONTRACT_CONTEXTS, 8960), CONTRACT_CONTEXT_VALUES + ) + + def test_rejects_missing_or_malformed_context_witness(self) -> None: + del self.receipt["contexts"]["4097"] + self._expect_rejection() + self.receipt["contexts"]["4097"] = _witness(8) + self.receipt["contexts"]["4097"]["logits_sha256"] = "0" * 63 + self._expect_rejection() + self.receipt["contexts"]["4097"] = _witness(8) + self.receipt["contexts"]["4097"]["prefill_token"] = K2_VOCAB_SIZE + self._expect_rejection() + self.receipt["contexts"]["4097"] = _witness(8) + del self.receipt["contexts"]["4097"]["layer0_qk_sha256"] + self._expect_rejection() + self.receipt["contexts"] = [] + self._expect_rejection() + + +class ProductionOracleBindingTest(unittest.TestCase): + def setUp(self) -> None: + directory = tempfile.TemporaryDirectory() + self.addCleanup(directory.cleanup) + self.root = Path(directory.name) + (self.root / "assistant").mkdir() + (self.root / "target").mkdir() + (self.root / "target" / "generation_config.json").write_text( + json.dumps({"eos_token_id": STOP_TOKENS}), encoding="utf-8" + ) + self.manifest_path = self.root / "mtp-manifest.json" + self.receipt_path = self.root / "target_prefill_oracle.json" + self.output = self.root / "mtp_spec_oracle.json" + self.manifest: dict[str, Any] = { + "abi": copy.deepcopy(K2_ROUND_ABI), + "artifacts": [ + {"path": "model.pte", "role": "pte"}, + *[{"path": f"model{index}.ptd", "role": "ptd"} for index in range(3)], + ], + "checkpoints": {"assistant": "assistant", "target": "target"}, + "max_context_length": 8960, + "method": "k2_round", + "ptd_order": [f"model{index}.ptd" for index in range(3)], + "schema_version": 1, + } + self.receipt: dict[str, Any] = { + "authority": "portable_eager", + "contexts": { + str(context): _witness(index) + for index, context in enumerate(CONTRACT_CONTEXT_VALUES) + }, + "schema_version": 1, + } + self._write() + self.combined_root = self.root / "combined" + receipts = self.combined_root / "receipts" + receipts.mkdir(parents=True) + self.production_mtp_path = receipts / "mtp.json" + self.production_mtp: dict[str, Any] = { + "acquisition": { + "assistant": ASSISTANT_CHECKPOINT_ACQUISITION, + "target": CHECKPOINT_ACQUISITION, + }, + "artifacts": copy.deepcopy(self.manifest["artifacts"]), + "export": {"max_seq_len": 8960, "methods": ["k2_round"]}, + "provenance": MTP_SOURCE_VERIFIED_PROVENANCE, + "ptd_order": copy.deepcopy(self.manifest["ptd_order"]), + } + self.production_mtp_path.write_text( + json.dumps(self.production_mtp), encoding="utf-8" + ) + self.production_receipt_path = receipts / "target_prefill.json" + self.production_receipt: dict[str, Any] = { + "authority": TARGET_PREFILL_AUTHORITY, + "checkpoint_acquisition": CHECKPOINT_ACQUISITION, + "contexts": { + str(context): _target_prefill_v2_context(context, index) + for index, context in enumerate(TARGET_PREFILL_CONTEXTS) + }, + "envelope_kind": TARGET_PREFILL_ENVELOPE_KIND, + "producer": { + "fbsource_commit": "1" * 40, + "runtime_source_receipt": {"bytes": 10, "sha256": "2" * 64}, + "source_path": reviewed_producer_source_path().name, + "source_sha256": hashlib.sha256( + reviewed_producer_source_path().read_bytes() + ).hexdigest(), + }, + "run": { + "command": ["owner-run"], + "finished_at_utc": "2026-08-07T01:00:00Z", + "host": "owner-host", + "started_at_utc": "2026-08-07T00:00:00Z", + }, + "schema_version": TARGET_PREFILL_SCHEMA_VERSION, + } + self._write_production_receipt() + self.combined_envelope_path = ( + self.combined_root / "gemma4_webgpu_combined_runtime.json" + ) + self.combined_envelope: dict[str, Any] = { + "receipts": { + "mtp": {"path": "receipts/mtp.json", "root": "mtp"}, + "target_prefill": {"path": "receipts/target_prefill.json"}, + }, + "schema_version": 3, + "source_verification": { + "mtp": {"provenance": MTP_SOURCE_VERIFIED_PROVENANCE} + }, + } + self.combined_envelope_path.write_text( + json.dumps(self.combined_envelope), encoding="utf-8" + ) + + def _write(self) -> None: + self.manifest_path.write_text(json.dumps(self.manifest), encoding="utf-8") + self.receipt_path.write_text(json.dumps(self.receipt), encoding="utf-8") + + def _build(self) -> dict[str, Any]: + return build_oracle_binding( + self.manifest_path, + self.receipt_path, + "portable_eager", + CONTRACT_CONTEXTS, + self.output, + ) + + def _records(self, binding: dict[str, Any]) -> dict[str, dict[str, Any]]: + records: dict[str, dict[str, Any]] = {} + for context in binding["contexts"]: + key = str(context) + raw = { + "bonus": 13, + "candidates": [11, 12], + "match_count": 2, + "state_probe": 0.0, + "target_greedy": [11, 12, 13], + } + decision = reconcile_k2_round( + raw["candidates"], + raw["target_greedy"], + raw["match_count"], + raw["bonus"], + raw["state_probe"], + context, + ORACLE_TOKEN_BUDGET - 1, + binding["stop_tokens"], + ) + round_record = {"kv_witness": "a" * 64, **raw, **decision} + target_prefill = binding["target_prefill"][key] + prefill_token = target_prefill.get( + "prefill_token_raw", target_prefill.get("prefill_token") + ) + records[key] = { + "accepted_prefix": [decision["accepted_drafts"]], + "bonus_accounting": [decision["next_seed"]], + "decoded_text": f"context {key}", + "kv_witnesses": [round_record["kv_witness"]], + "reset_replay": raw, + "rounds": [round_record], + "selected_logits": [decision["selected"]], + "stop_handling": { + "discarded": [decision["discarded"]], + "stop_token": None, + }, + "target_prefill": target_prefill, + "useful_tokens": [ + prefill_token, + *decision["committed"], + ], + } + return records + + def _write_production_receipt(self) -> None: + self.production_receipt_path.write_bytes( + canonical_json_bytes(self.production_receipt) + ) + + def _build_production(self, contexts: str = CONTRACT_CONTEXTS) -> dict[str, Any]: + with mock.patch( + "executorch.examples.models.gemma4.tests.generate_mtp_spec_oracle." + "validate_combined_runtime_envelope" + ) as validator: + binding = build_oracle_binding( + self.manifest_path, + None, + "portable_eager", + contexts, + self.output, + combined_runtime_root=self.combined_root, + ) + validator.assert_called_once_with(self.combined_root, self.combined_envelope) + return binding + + def _build_production_direct(self) -> dict[str, Any]: + with mock.patch( + "executorch.examples.models.gemma4.tests.generate_mtp_spec_oracle." + "validate_combined_runtime_envelope" + ) as validator: + binding = mtp_oracle.build_production_oracle_binding( + self.production_mtp_path, + self.root / "target", + self.root / "assistant", + "portable_eager", + CONTRACT_CONTEXTS, + self.output, + combined_runtime_root=self.combined_root, + ) + validator.assert_called_once_with(self.combined_root, self.combined_envelope) + return binding + + def test_production_binding_derives_staged_receipt_and_is_acceptable(self) -> None: + binding = self._build_production_direct() + document = assemble_oracle_document(binding, self._records(binding)) + self.assertEqual(document["schema_version"], 2) + self.assertEqual(document["closure_state"], "full") + self.assertEqual(document["target_prefill_authority"], "bound") + self.assertEqual(document["replay_independence"], "eager_vs_lowered_only") + self.assertTrue(production_oracle_is_acceptable(document)) + self.assertEqual(document["contexts"], list(TARGET_PREFILL_CONTEXTS)) + self.assertEqual( + binding["target_prefill"]["512"]["prefill_token_raw"], + self.production_receipt["contexts"]["512"]["prefill_token_raw"], + ) + self.assertEqual( + document["mtp_manifest_sha256"], + hashlib.sha256(self.production_mtp_path.read_bytes()).hexdigest(), + ) + + def test_legacy_receipt_stays_non_accepting(self) -> None: + binding = self._build() + document = assemble_oracle_document(binding, self._records(binding)) + self.assertEqual(document["schema_version"], 1) + self.assertEqual(document["closure_state"], "absent") + self.assertEqual(document["target_prefill_authority"], "legacy_unbound") + self.assertFalse(production_oracle_is_acceptable(document)) + + def test_legacy_combined_binding_uses_the_staged_mtp_identity(self) -> None: + binding = self._build_production() + document = assemble_oracle_document(binding, self._records(binding)) + self.assertEqual( + hashlib.sha256(self.production_mtp_path.read_bytes()).hexdigest(), + document["mtp_manifest_sha256"], + ) + self.assertTrue(production_oracle_is_acceptable(document)) + + def test_production_binding_requires_exact_ten_contexts(self) -> None: + with self.assertRaisesRegex(OracleError, "exact ten contexts"): + self._build_production("128,512") + del self.production_receipt["contexts"]["4097"] + self._write_production_receipt() + with self.assertRaisesRegex(OracleError, "exact ten contexts"): + self._build_production() + + def test_full_closure_requires_source_verified_mtp_provenance(self) -> None: + self.production_mtp_path.write_text( + json.dumps( + { + "provenance": { + "artifact_status": "accepted_behavior_oracle", + "source_closure": "pending_final_source_rebuild", + } + } + ), + encoding="utf-8", + ) + with self.assertRaisesRegex(OracleError, "not source verified"): + self._build_production() + + def test_target_receipt_authority_domain_is_separate_from_replay(self) -> None: + self.assertEqual(SUPPORTED_AUTHORITIES, ("portable_eager",)) + self.assertEqual(TARGET_PREFILL_AUTHORITIES, ("target_only_eager",)) + self.production_receipt["authority"] = "portable_eager" + self._write_production_receipt() + with self.assertRaisesRegex(OracleError, "target-prefill authority"): + self._build_production() + + def test_raw_post_token_mismatch_and_k2_bonus_mismatch_fail_closed(self) -> None: + witness = self.production_receipt["contexts"]["512"] + witness["prefill_token_post_softcap"] += 1 + self._write_production_receipt() + with self.assertRaisesRegex(OracleError, "raw/post-softcap"): + self._build_production() + witness["prefill_token_post_softcap"] = witness["prefill_token_raw"] + require_prefill_token_match(512, witness["prefill_token_raw"], witness) + with self.assertRaisesRegex(OracleError, "prefill token"): + require_prefill_token_match(512, witness["prefill_token_raw"] + 1, witness) + + def test_production_validator_value_error_is_translated(self) -> None: + with mock.patch( + "executorch.examples.models.gemma4.tests.generate_mtp_spec_oracle." + "validate_combined_runtime_envelope", + side_effect=ValueError("source closure rejected"), + ): + with self.assertRaisesRegex(OracleError, "source closure rejected"): + build_oracle_binding( + self.manifest_path, + None, + "portable_eager", + CONTRACT_CONTEXTS, + self.output, + combined_runtime_root=self.combined_root, + ) + + def test_acceptance_rejects_unknown_or_incomplete_states(self) -> None: + binding = self._build_production_direct() + document = assemble_oracle_document(binding, self._records(binding)) + for key, value in ( + ("closure_state", "absent"), + ("target_prefill_authority", "legacy_unbound"), + ): + with self.subTest(key=key): + mutated = copy.deepcopy(document) + mutated[key] = value + self.assertFalse(production_oracle_is_acceptable(mutated)) + mutated = copy.deepcopy(document) + mutated["closure_state"] = "artifact_validated_pending_source" + with self.assertRaisesRegex(OracleError, "closure_state"): + production_oracle_is_acceptable(mutated) + mutated = copy.deepcopy(document) + del mutated["production_binding"] + self.assertFalse(production_oracle_is_acceptable(mutated)) + + def test_acceptance_rejects_malformed_oracle_content(self) -> None: + binding = self._build_production_direct() + document = assemble_oracle_document(binding, self._records(binding)) + mutations = ( + lambda value: value.__setitem__("abi", {}), + lambda value: value.__setitem__("method", "forward"), + lambda value: value.__setitem__("authority", "target_only_eager"), + lambda value: value.__setitem__("records", {}), + lambda value: value["records"]["512"].__setitem__("rounds", []), + lambda value: value["records"]["512"].__setitem__( + "reset_replay", {} + ), + lambda value: value["production_binding"].__setitem__( + "mtp_manifest_sha256", "0" * 64 + ), + ) + for mutate in mutations: + with self.subTest(mutation=mutate): + mutated = copy.deepcopy(document) + mutate(mutated) + self.assertFalse(production_oracle_is_acceptable(mutated)) + + def test_acceptance_rejects_target_prefill_receipt_content_drift(self) -> None: + binding = self._build_production_direct() + document = assemble_oracle_document(binding, self._records(binding)) + mutated = copy.deepcopy(document) + mutated["records"]["512"]["target_prefill"]["logits_pre_softcap"][ + "sha256" + ] = "0" * 64 + self.assertFalse(production_oracle_is_acceptable(mutated)) + + def test_main_refuses_to_write_an_unacceptable_production_oracle(self) -> None: + binding = self._build_production_direct() + records = self._records(binding) + records["512"]["target_prefill"]["logits_pre_softcap"]["sha256"] = ( + "0" * 64 + ) + with mock.patch.object( + mtp_oracle, + "build_production_oracle_binding", + return_value=binding, + ), mock.patch.object( + mtp_oracle, + "run_portable_eager_authority", + return_value=records, + ): + with self.assertRaisesRegex(OracleError, "failed production validation"): + mtp_oracle.main( + [ + "--production-mtp-manifest", + str(self.production_mtp_path), + "--target-checkpoint", + str(self.root / "target"), + "--assistant-checkpoint", + str(self.root / "assistant"), + "--combined-runtime-root", + str(self.combined_root), + "--authority", + "portable_eager", + "--contexts", + CONTRACT_CONTEXTS, + "--output", + str(self.output), + ] + ) + self.assertFalse(self.output.exists()) + + def test_cli_requires_combined_root_and_has_no_loose_receipt_flag(self) -> None: + parser = _build_parser() + help_text = " ".join(parser.format_help().split()) + option_strings = { + option for action in parser._actions for option in action.option_strings + } + self.assertIn("target-only eager prefill evidence", help_text) + self.assertIn("does not independently validate the shared model", help_text) + self.assertIn("--oracle-binding-manifest", option_strings) + self.assertIn("--production-mtp-manifest", option_strings) + self.assertIn("--target-checkpoint", option_strings) + self.assertIn("--assistant-checkpoint", option_strings) + self.assertIn("--combined-runtime-root", option_strings) + self.assertNotIn("--mtp-manifest", option_strings) + self.assertNotIn("--target-prefill-oracle", option_strings) + with self.assertRaises(SystemExit): + parser.parse_args( + [ + "--oracle-binding-manifest", + str(self.manifest_path), + "--authority", + "portable_eager", + "--contexts", + CONTRACT_CONTEXTS, + "--output", + str(self.output), + ] + ) + + def test_k2_abi_counts_match_production_edge_census(self) -> None: + validate_k2_abi_edge_census(MTP_EDGE_CENSUS) + mutated = dict(MTP_EDGE_CENSUS) + mutated["topk"] = 1 + with self.assertRaisesRegex(OracleError, "K=2 ABI"): + validate_k2_abi_edge_census(mutated) + + def test_d6_and_d8_tensor_encoders_match(self) -> None: + tensors = ( + torch.arange(12, dtype=torch.float32).reshape(3, 4), + torch.arange(12, dtype=torch.float32).reshape(3, 4).transpose(0, 1), + torch.arange(12, dtype=torch.int64).reshape(3, 4), + torch.arange(12, dtype=torch.int32).reshape(3, 4).transpose(0, 1), + ) + for tensor in tensors: + with self.subTest(dtype=tensor.dtype, contiguous=tensor.is_contiguous()): + self.assertEqual( + hashlib.sha256(target_prefill_tensor_bytes(tensor)).hexdigest(), + assistant_tensor_sha256(tensor), + ) + + +class K2ReconcilerTest(unittest.TestCase): + def _round( + self, + candidates: list[int], + greedy: list[int], + match_count: int, + bonus: int, + start_position: int = 2, + token_budget: int = 3, + stop_tokens: list[int] | None = None, + state_probe: float = 0.0, + vocab_size: int = K2_VOCAB_SIZE, + ) -> dict[str, Any]: + return reconcile_k2_round( + candidates, + greedy, + match_count, + bonus, + state_probe, + start_position, + token_budget, + STOP_TOKENS if stop_tokens is None else stop_tokens, + vocab_size, + ) + + def test_match_counts_advance_position_by_match_plus_one(self) -> None: + expected = ((0, [90], 3), (1, [10, 90], 4), (2, [10, 11, 90], 5)) + greedy_by_match = { + 0: [90, 91, 92], + 1: [10, 90, 92], + 2: [10, 11, 90], + } + for match_count, committed, next_position in expected: + with self.subTest(match_count=match_count): + decision = self._round( + [10, 11], greedy_by_match[match_count], match_count, 90 + ) + self.assertTrue(decision["valid"]) + self.assertEqual(decision["committed"], committed) + self.assertEqual(decision["selected"], committed) + self.assertEqual(decision["next_position"], next_position) + self.assertEqual(decision["next_seed"], 90) + self.assertEqual(decision["accepted_drafts"], match_count) + self.assertFalse(decision["stopped"]) + self.assertEqual(decision["discarded"], []) + + def test_chained_rounds_walk_start_positions_two_three_five(self) -> None: + first = self._round([10, 11], [90, 91, 92], 0, 90, start_position=2) + self.assertEqual(first["next_position"], 3) + self.assertEqual(first["next_seed"], 90) + second = self._round([20, 21], [20, 91, 92], 1, 91, start_position=3) + self.assertEqual(second["next_position"], 5) + self.assertEqual(second["next_seed"], 91) + third = self._round([30, 31], [30, 31, 92], 2, 92, start_position=5) + self.assertEqual(third["next_position"], 8) + self.assertEqual(third["next_seed"], 92) + + def test_seeds_from_bonus_never_from_draft_or_target_tail(self) -> None: + decision = self._round([10, 11], [10, 90, 92], 1, 90) + self.assertEqual(decision["next_seed"], 90) + self.assertNotEqual(decision["next_seed"], 10) + self.assertNotEqual(decision["next_seed"], 92) + self.assertEqual(decision["selected"][-1], decision["next_seed"]) + + def test_stop_token_is_not_committed_and_remainder_discarded(self) -> None: + decision = self._round([106, 11], [106, 11, 90], 2, 90) + self.assertTrue(decision["valid"]) + self.assertTrue(decision["stopped"]) + self.assertEqual(decision["stop_token"], 106) + self.assertEqual(decision["committed"], []) + self.assertEqual(decision["discarded"], [11, 90]) + self.assertEqual(decision["next_position"], 5) + + def test_stop_token_in_the_bonus_slot_commits_the_prefix(self) -> None: + decision = self._round([10, 11], [10, 11, 1], 2, 1) + self.assertTrue(decision["stopped"]) + self.assertEqual(decision["stop_token"], 1) + self.assertEqual(decision["committed"], [10, 11]) + self.assertEqual(decision["discarded"], []) + + def test_budget_truncation_discards_from_the_overflow_token(self) -> None: + decision = self._round([10, 11], [10, 11, 90], 2, 90, token_budget=1) + self.assertTrue(decision["valid"]) + self.assertFalse(decision["stopped"]) + self.assertEqual(decision["committed"], [10]) + self.assertEqual(decision["discarded"], [11, 90]) + decision = self._round([10, 11], [10, 11, 90], 2, 90, token_budget=2) + self.assertEqual(decision["committed"], [10, 11]) + self.assertEqual(decision["discarded"], [90]) + decision = self._round([10, 11], [10, 11, 90], 2, 90, token_budget=3) + self.assertEqual(decision["committed"], [10, 11, 90]) + self.assertEqual(decision["discarded"], []) + + def test_rejects_self_inconsistent_match_metadata(self) -> None: + self.assertFalse(self._round([10, 11], [10, 91, 92], 2, 92)["valid"]) + self.assertFalse(self._round([10, 11], [10, 11, 90], 1, 11)["valid"]) + self.assertFalse(self._round([10, 11], [90, 91, 92], 1, 91)["valid"]) + self.assertFalse(self._round([10, 11], [10, 90, 92], 1, 92)["valid"]) + self.assertFalse(self._round([10, 11], [10, 11, 90], 2, 11)["valid"]) + + def test_rejects_every_input_guard(self) -> None: + good = ([10, 11], [10, 11, 90], 2, 90) + self.assertTrue(self._round(*good)["valid"]) + self.assertFalse(self._round(*good, start_position=1)["valid"]) + self.assertFalse(self._round(*good, start_position=0)["valid"]) + self.assertFalse(self._round(*good, start_position=-1)["valid"]) + self.assertFalse(self._round(*good, token_budget=0)["valid"]) + self.assertFalse(self._round(*good, vocab_size=0)["valid"]) + self.assertFalse(self._round(*good, vocab_size=-1)["valid"]) + self.assertFalse(self._round([10, 11], [10, 11, 90], -1, 90)["valid"]) + self.assertFalse(self._round([10, 11], [10, 11, 90], 3, 90)["valid"]) + for probe in (float("nan"), float("inf"), float("-inf")): + with self.subTest(probe=probe): + self.assertFalse(self._round(*good, state_probe=probe)["valid"]) + self.assertTrue(self._round(*good, start_position=2)["valid"]) + + def test_rejects_out_of_range_tokens(self) -> None: + self.assertFalse(self._round([-1, 11], [10, 11, 90], 0, 10)["valid"]) + self.assertFalse(self._round([K2_VOCAB_SIZE, 11], [10, 11, 90], 0, 10)["valid"]) + self.assertFalse(self._round([10, 11], [-1, 11, 90], 0, 90)["valid"]) + self.assertFalse( + self._round([10, 11], [10, 11, K2_VOCAB_SIZE], 2, K2_VOCAB_SIZE)["valid"] + ) + self.assertFalse( + self._round([10, 11], [10, 11, 90], 2, 90, vocab_size=90)["valid"] + ) + self.assertTrue( + self._round([10, 11], [10, 11, 90], 2, 90, vocab_size=91)["valid"] + ) + + def test_rejected_decision_uses_documented_defaults(self) -> None: + decision = self._round([10, 11], [10, 91, 92], 2, 92) + self.assertEqual( + decision, + { + "accepted_drafts": 0, + "committed": [], + "discarded": [], + "next_position": -1, + "next_seed": -1, + "selected": [], + "stop_token": -1, + "stopped": False, + "valid": False, + }, + ) + + def test_empty_stop_token_list_never_stops(self) -> None: + decision = self._round([106, 1], [106, 1, 90], 2, 90, stop_tokens=[]) + self.assertTrue(decision["valid"]) + self.assertFalse(decision["stopped"]) + self.assertEqual(decision["stop_token"], -1) + self.assertEqual(decision["committed"], [106, 1, 90]) + + +GEMMA4_ANCHOR = "examples/models/gemma4/targets.bzl" +SPEC_RUNNER_HEADER = "examples/models/gemma4/runner/gemma4_spec_runner.h" +SOURCE_ROOT_ENV = "EXECUTORCH_SOURCE_ROOT" + +INPUT_GUARD_ANCHOR = "start_position <" +SELF_CONSISTENCY_ANCHOR = "expected_matches" + +VALID_ROUND: dict[str, Any] = { + "candidates": [10, 11], + "target_greedy": [10, 11, 90], + "match_count": 2, + "bonus": 90, + "state_probe": 0.0, + "start_position": 2, + "token_budget": 3, + "stop_tokens": STOP_TOKENS, + "vocab_size": K2_VOCAB_SIZE, +} +INPUT_GUARD_PROBES: dict[str, dict[str, Any]] = { + "start_position < 2": {"start_position": 1}, + "token_budget == 0": {"token_budget": 0}, + "vocab_size <= 0": {"vocab_size": 0}, + "match_count < 0": {"match_count": -1}, + "match_count > 2": {"match_count": 3}, + "!std::isfinite(state_probe)": {"state_probe": float("nan")}, +} +SELF_CONSISTENCY_PROBES: dict[str, dict[str, Any]] = { + "match_count != expected_matches": {"match_count": 1, "bonus": 11}, + "!valid_token(bonus)": {"bonus": -1}, + "bonus != target_greedy[match_count]": {"bonus": 91}, +} + +_ASSIGNMENT_PATTERN: re.Pattern[str] = re.compile(r"decision\.(\w+)\s*=\s*([^;]+);") +_GUARD_PATTERN: re.Pattern[str] = re.compile(r"if\s*\(([^;{}]*?)\)\s*\{", re.S) + + +def _root_candidates() -> list[tuple[str, Path | None]]: + override = os.environ.get(SOURCE_ROOT_ENV) + try: + package = importlib.util.find_spec("executorch") + except (ImportError, ValueError): + package = None + staged = list(package.submodule_search_locations or ()) if package else [] + here = Path(__file__).resolve() + walked = next( + (parent for parent in here.parents if (parent / GEMMA4_ANCHOR).is_file()), None + ) + return [ + (f"${SOURCE_ROOT_ENV}", Path(override) if override else None), + ("`executorch` package runfile", Path(staged[0]) if staged else None), + (f"__file__ walk above {here}", walked), + ] + + +def _source_root() -> Path: + attempted: list[str] = [] + for strategy, candidate in _root_candidates(): + attempted.append(f"{strategy} -> {candidate}") + if candidate is not None and (candidate / GEMMA4_ANCHOR).is_file(): + return candidate + raise FileNotFoundError( + f"no ExecuTorch source root containing {GEMMA4_ANCHOR}; " + f"tried {'; '.join(attempted)}" + ) + + +def _read_spec_runner_header() -> str: + path = _source_root() / SPEC_RUNNER_HEADER + if not path.is_file(): + raise FileNotFoundError( + f"missing D9 source under test: {path}; `reconcile_k2_round` cannot be " + "cross-checked against `reconcile_gemma4_k2` until it lands" + ) + return path.read_text(encoding="utf-8") + + +def _normalize(expression: str) -> str: + return " ".join(expression.replace("output.", "").split()) + + +def _cpp_assignment(header: str, field: str) -> str: + found = [ + match.group(2) + for match in _ASSIGNMENT_PATTERN.finditer(header) + if match.group(1) == field and re.search(r"[A-Za-z_]", match.group(2)) + ] + if len(found) != 1: + raise AssertionError( + f"expected exactly one computed `decision.{field} = ...;` in " + f"{SPEC_RUNNER_HEADER}, found {len(found)}" + ) + return _normalize(found[0]) + + +def _cpp_guard_clauses(header: str, anchor: str) -> set[str]: + found = [ + match.group(1) + for match in _GUARD_PATTERN.finditer(header) + if anchor in match.group(1) + ] + if len(found) != 1: + raise AssertionError( + f"expected exactly one `if` clause naming {anchor!r} in " + f"{SPEC_RUNNER_HEADER}, found {len(found)}" + ) + return {_normalize(clause) for clause in found[0].split("||")} + + +def _mirror_expression(field: str) -> str: + tree = ast.parse(inspect.getsource(reconcile_k2_round)) + found = [ + ast.unparse(value) + for node in ast.walk(tree) + if isinstance(node, ast.Dict) + for key, value in zip(node.keys, node.values) + if isinstance(key, ast.Constant) + and key.value == field + and any(isinstance(child, ast.Name) for child in ast.walk(value)) + ] + if len(found) != 1: + raise AssertionError( + f"expected exactly one computed {field!r} entry in " + f"reconcile_k2_round, found {len(found)}" + ) + return found[0] + + +class K2MirrorMatchesSpecRunnerHeaderTest(unittest.TestCase): + """Pins `reconcile_k2_round` to `reconcile_gemma4_k2` in the D9 header.""" + + header: str + + def setUp(self) -> None: + self.header = _read_spec_runner_header() + + def _assert_probes_are_rejected( + self, probes: Mapping[str, Mapping[str, Any]] + ) -> None: + self.assertTrue(reconcile_k2_round(**VALID_ROUND)["valid"]) + for clause, probe in probes.items(): + with self.subTest(clause=clause): + self.assertFalse( + reconcile_k2_round(**{**VALID_ROUND, **probe})["valid"] + ) + + def test_progression_and_seed_expressions_match_the_mirror(self) -> None: + for field in ("next_position", "next_seed"): + with self.subTest(field=field): + self.assertEqual( + _cpp_assignment(self.header, field), _mirror_expression(field) + ) + + def test_every_input_guard_clause_is_enforced_by_the_mirror(self) -> None: + self.assertEqual( + _cpp_guard_clauses(self.header, INPUT_GUARD_ANCHOR), + set(INPUT_GUARD_PROBES), + ) + self._assert_probes_are_rejected(INPUT_GUARD_PROBES) + + def test_every_self_consistency_clause_is_enforced_by_the_mirror(self) -> None: + self.assertEqual( + _cpp_guard_clauses(self.header, SELF_CONSISTENCY_ANCHOR), + set(SELF_CONSISTENCY_PROBES), + ) + self._assert_probes_are_rejected(SELF_CONSISTENCY_PROBES) diff --git a/examples/models/gemma4/tests/test_oss_source_closure.py b/examples/models/gemma4/tests/test_oss_source_closure.py new file mode 100644 index 00000000000..5128c82b881 --- /dev/null +++ b/examples/models/gemma4/tests/test_oss_source_closure.py @@ -0,0 +1,688 @@ +# 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. + +# pyre-strict + +"""D10 OSS source closure and plain-Gemma non-disturbance contract.""" + +from __future__ import annotations + +import ast +import hashlib +import importlib.util +import json +import re +import shutil +import subprocess +import sys +import unittest + +from pathlib import Path +from types import ModuleType + +_XPLAT_PREFIX = ("xplat", "executorch") +_FBCODE_PREFIX = ("fbcode", "executorch") + +_SL_STATUS_REASON = "list D10 changed paths for the OSS closure gate - sl help status" +_SL_LOG_REASON = "find the commit that introduced this gate - sl help log" + +_D10_SURFACE_DIRECTORIES = ( + "backends/webgpu/scripts", + "backends/webgpu/test", + "examples/models/gemma4", +) + +_D10_SURFACE_FILES = ("backends/webgpu/CMakeLists.txt",) + +_THIS_GATE = "examples/models/gemma4/tests/test_oss_source_closure.py" + +# D10's changed files, executorch-relative. The drift guard below compares this +# list against source control wherever source control can still identify D10. +_D10_CHANGED_PATHS = ( + "backends/webgpu/CMakeLists.txt", + "backends/webgpu/scripts/test_gemma4_wasm_factory_contract.sh", + "backends/webgpu/scripts/test_webgpu_native_ci.sh", + "backends/webgpu/test/BUCK", + "backends/webgpu/test/native/test_q4gsw_m3.cpp", + "backends/webgpu/test/native/test_scatter.cpp", + "backends/webgpu/test/native/test_topk.cpp", + "backends/webgpu/test/op_tests/cases.py", + "backends/webgpu/test/op_tests/test_typed_input_contract.py", + "backends/webgpu/test/ops/index/test_index.py", + "backends/webgpu/test/ops/scatter/__init__.py", + "backends/webgpu/test/ops/scatter/export_scatter_artifacts.py", + "backends/webgpu/test/ops/scatter/test_scatter.py", + "backends/webgpu/test/ops/test_gather.py", + "backends/webgpu/test/ops/test_to_copy.py", + "backends/webgpu/test/ops/test_where.py", + "backends/webgpu/test/ops/topk/__init__.py", + "backends/webgpu/test/ops/topk/export_topk_artifacts.py", + "backends/webgpu/test/ops/topk/test_topk.py", + "backends/webgpu/test/targets.bzl", + "backends/webgpu/test/test_native_ci_contract.py", + "examples/models/gemma4/tests/generate_mtp_spec_oracle.py", + "examples/models/gemma4/tests/targets.bzl", + "examples/models/gemma4/tests/test_eagle_combined_round.py", + "examples/models/gemma4/tests/test_export_assistant_webgpu_artifacts.py", + "examples/models/gemma4/tests/test_export_partitioners.py", + "examples/models/gemma4/tests/test_gemma4_spec_runner_contract.cpp", + "examples/models/gemma4/tests/test_mtp_spec_oracle.py", + "examples/models/gemma4/tests/test_oss_source_closure.py", + "examples/models/gemma4/tests/test_webgpu_artifact_manifest.py", + "examples/models/gemma4/tests/test_webgpu_spec_contract.py", +) + +_D10_COMMAND_MODULES = ( + "executorch.backends.webgpu.test.ops.scatter.export_scatter_artifacts", + "executorch.backends.webgpu.test.ops.topk.export_topk_artifacts", + "executorch.backends.webgpu.test.ops.topk.test_topk", + "executorch.backends.webgpu.test.op_tests.generate_op_tests", +) + +_D10_COMMAND_SCRIPTS = ("examples/models/gemma4/tests/generate_mtp_spec_oracle.py",) + +_D10_COMMAND_TEST_CASE = ( + "executorch.backends.webgpu.test.ops.topk.test_topk", + "TestEagleTopKCpu", + "test_eager_reference_is_repeatable", +) + +# Each pattern splits one character so this file is not its own violation. +_INTERNAL_REFERENCE_PATTERNS = ( + r"manifol[d]", + r"internalf[b]", + r"fbur[l]", + r"/data/user[s]/", + r"/hom[e]/", + r"/User[s]/", + r"/mn[t]/", + r"examples/models/f[b]/", + r"\bD1[0-9]{8}\b", + r"localhos[t]:[0-9]+", + r"127[.]0[.]0[.]1:[0-9]+", + r"\.intern[.]facebook[.]com", +) + +# Rejection tests must name the strings they reject; only these two may do so. +_NEGATIVE_FIXTURE_MARKER = "oss-closure-" "fixture" # split: not self-marking +_NEGATIVE_FIXTURE_FILES = ( + "examples/models/gemma4/tests/test_export_assistant_webgpu_artifacts.py", + "examples/models/gemma4/tests/test_webgpu_artifact_manifest.py", +) + +_BINARY_ARTIFACT_SUFFIXES = ( + ".pte", + ".ptd", + ".bin", + ".gguf", + ".safetensors", + ".png", + ".wasm", +) + +_PLAIN_PARTITIONER = "examples/models/gemma4/webgpu_partitioner.py" +_PLAIN_MANIFEST_MODULE = "examples/models/gemma4/webgpu_artifact_manifest.py" +_PLAIN_MANIFEST_JSON = "examples/models/gemma4/manifests/gemma4_e2b_webgpu.json" +_PLAIN_RUNNER_HEADER = "examples/models/gemma4/runner/gemma4_runner.h" +_PLAIN_RUNNER_SOURCE = "examples/models/gemma4/runner/gemma4_runner.cpp" +_PLAIN_MODEL_TARGETS = "examples/models/gemma4/targets.bzl" + +_PLAIN_WEBGPU_ALLOWLIST = ( + "exir_ops.edge.aten._assert_scalar.default", + "exir_ops.edge.aten.add.Tensor", + "exir_ops.edge.aten.argmax.default", + "exir_ops.edge.aten.cat.default", + "exir_ops.edge.aten.clamp.default", + "exir_ops.edge.aten.div.Tensor", + "exir_ops.edge.aten.gelu.default", + "exir_ops.edge.aten.mul.Tensor", + "exir_ops.edge.aten.permute_copy.default", + "exir_ops.edge.aten.select_copy.int", + "exir_ops.edge.aten.sigmoid.default", + "exir_ops.edge.aten.slice_copy.Tensor", + "exir_ops.edge.aten.squeeze_copy.dims", + "exir_ops.edge.aten.sym_constrain_range_for_size.default", + "exir_ops.edge.aten.tanh.default", + "exir_ops.edge.aten.unsqueeze_copy.default", + "exir_ops.edge.aten.view_copy.default", + "exir_ops.edge.dim_order_ops._clone_dim_order.default", + "exir_ops.edge.dim_order_ops._to_dim_order_copy.default", + "exir_ops.edge.et_vk.apply_rotary_emb_hf.default", + "exir_ops.edge.et_vk.apply_rotary_emb_hf_single.default", + "exir_ops.edge.et_vk.gemma4_sdpa.default", + "exir_ops.edge.et_vk.rms_norm.default", + "exir_ops.edge.et_vk.select_as_symint.default", +) + +_PLAIN_EXPORT_CONTRACT_SHA256 = ( + "7b1e7ad9753bdc937f32a254eff04dad10528ac3c704f0f176894c2c435c9f2d" +) +_PLAIN_ARCHITECTURE_SHA256 = ( + "d731d17637aca0e808a6bdca6b80231310b84b6ed1708c579d95ff96c470d1a8" +) +_PLAIN_ACQUISITION_SHA256 = ( + "f1aad2baf1b48edf5124b993208f73ac2ae6878d890a84ab0e211d05babf316a" +) + + +def _tree_root() -> tuple[str, Path]: + parents = list(Path(__file__).resolve().parents) + for parent in parents: + if (parent.joinpath(*_XPLAT_PREFIX).is_dir()) and ( + parent.joinpath(*_FBCODE_PREFIX).is_dir() + ): + return "fbsource", parent + for parent in parents: + if (parent / "backends" / "webgpu").is_dir() and ( + parent / "examples" / "models" / "gemma4" + ).is_dir(): + return "oss", parent + raise RuntimeError(f"no fbsource or OSS root above {Path(__file__).resolve()}") + + +def _physical_paths(layout: str, root: Path, relative: str) -> dict[str, Path]: + if layout == "oss": + return {"oss": root / relative} + return { + "xplat": root.joinpath(*_XPLAT_PREFIX) / relative, + "fbcode": root.joinpath(*_FBCODE_PREFIX) / relative, + } + + +def _canonical_paths(layout: str, root: Path, relative: str) -> Path: + forms = _physical_paths(layout, root, relative) + return forms["oss"] if layout == "oss" else forms["xplat"] + + +def _identity_sha256(value: object) -> str: + return hashlib.sha256(json.dumps(value, sort_keys=True).encode("utf-8")).hexdigest() + + +def _dotted_name(node: ast.expr) -> str: + parts: list[str] = [] + while isinstance(node, ast.Attribute): + parts.append(node.attr) + node = node.value + if not isinstance(node, ast.Name): + raise ValueError(f"not an attribute chain: {ast.dump(node)}") + parts.append(node.id) + return ".".join(reversed(parts)) + + +def _module_tree(path: Path) -> ast.Module: + return ast.parse(path.read_text(encoding="utf-8"), filename=str(path)) + + +def _module_constant(tree: ast.Module, name: str) -> object: + for node in tree.body: + if isinstance(node, ast.Assign) and any( + isinstance(target, ast.Name) and target.id == name + for target in node.targets + ): + return ast.literal_eval(node.value) + raise AssertionError(f"module constant not found: {name}") + + +def _function_def(tree: ast.Module, name: str) -> ast.FunctionDef: + for node in tree.body: + if isinstance(node, ast.FunctionDef) and node.name == name: + return node + raise AssertionError(f"function not found: {name}") + + +def _method_def(tree: ast.Module, class_name: str, name: str) -> ast.FunctionDef: + for node in tree.body: + if isinstance(node, ast.ClassDef) and node.name == class_name: + for member in node.body: + if isinstance(member, ast.FunctionDef) and member.name == name: + return member + raise AssertionError(f"method not found: {class_name}.{name}") + + +def _returned_dotted_names(function: ast.FunctionDef) -> tuple[str, ...]: + for node in ast.walk(function): + if isinstance(node, ast.Return) and isinstance(node.value, ast.List): + return tuple(_dotted_name(element) for element in node.value.elts) + raise AssertionError(f"{function.name} does not return a list literal") + + +def _sl_lines(root: Path, args: tuple[str, ...], reason: str) -> tuple[str, ...] | None: + """Stdout lines of one read-only `sl` call; None when `sl` cannot answer.""" + try: + completed = subprocess.run( + ["sl", *args, "--reason", reason], + capture_output=True, + cwd=root, + text=True, + timeout=120, + ) + except (OSError, subprocess.SubprocessError): + return None + if completed.returncode != 0: + return None + return tuple(completed.stdout.splitlines()) + + +def _status_paths( + root: Path, codes: tuple[str, ...] = ("M ", "A ") +) -> frozenset[str] | None: + """Repo-relative paths `sl status` reports under the given status codes.""" + lines = _sl_lines(root, ("status",), _SL_STATUS_REASON) + if lines is None: + return None + return frozenset(line[2:] for line in lines if line[:2] in codes) + + +def _introducing_commit(root: Path) -> str: + """Node that last touched this gate; empty while D10 is uncommitted.""" + relative = str(Path(__file__).resolve().relative_to(root)) + lines = _sl_lines( + root, ("log", relative, "-T", "{node}\n", "-l", "1"), _SL_LOG_REASON + ) + return lines[0] if lines else "" + + +def _sapling_is_usable(root: Path) -> bool: + """True when `sl` and a Sapling working copy are both present at root.""" + if shutil.which("sl") is None: + return False + return any( + (parent / ".sl").is_dir() or (parent / ".hg").is_dir() + for parent in (root, *root.parents) + ) + + +def _executorch_relative(layout: str, repo_path: str) -> str | None: + if layout == "oss": + return repo_path + for parts in (_XPLAT_PREFIX, _FBCODE_PREFIX): + prefix = "/".join(parts) + "/" + if repo_path.startswith(prefix): + return repo_path[len(prefix) :] + return None + + +def _in_d10_surface(relative: str) -> bool: + if relative in _D10_SURFACE_FILES: + return True + return any( + relative.startswith(directory + "/") for directory in _D10_SURFACE_DIRECTORIES + ) + + +def _untracked_relatives(layout: str, root: Path) -> frozenset[str]: + """Executorch-relative paths source control reports as not tracked.""" + reported = _status_paths(root, ("? ",)) + if reported is None: + return frozenset() + relatives: set[str] = set() + for repo_path in reported: + relative = _executorch_relative(layout, repo_path) + if relative is not None: + relatives.add(relative) + return frozenset(relatives) + + +def _surface_status_paths(layout: str, root: Path) -> frozenset[str] | None: + """Reported changes narrowed to D10's surface; None when `sl` cannot answer.""" + reported = _status_paths(root) + if reported is None: + return None + relatives: set[str] = set() + for repo_path in reported: + relative = _executorch_relative(layout, repo_path) + if relative is not None and _in_d10_surface(relative): + relatives.add(relative) + return frozenset(relatives) + + +def _install_executorch_namespace(layout: str, root: Path) -> None: + """Bind `executorch.*` to this tree so a by-path load resolves its imports.""" + if "executorch" in sys.modules: + return + base = root / "xplat" / "executorch" if layout == "fbsource" else root + package = ModuleType("executorch") + package.__path__ = [str(base)] # pyre-ignore[16] + sys.modules["executorch"] = package + + +def _load_plain_manifest_module(layout: str, root: Path) -> ModuleType: + """Load by path; the gemma4 package __init__ pulls in torch.""" + _install_executorch_namespace(layout, root) + path = _canonical_paths(layout, root, _PLAIN_MANIFEST_MODULE) + spec = importlib.util.spec_from_file_location( + "gemma4_plain_webgpu_artifact_manifest", path + ) + if spec is None: + raise AssertionError(f"cannot load plain manifest module: {path}") + loader = spec.loader + if loader is None: + raise AssertionError(f"cannot load plain manifest module: {path}") + module = importlib.util.module_from_spec(spec) + loader.exec_module(module) + return module + + +class OssSourceClosureTest(unittest.TestCase): + maxDiff: int | None = None + + def setUp(self) -> None: + self.layout, self.root = _tree_root() + self.changed = _D10_CHANGED_PATHS + + def _existing_copies(self) -> list[tuple[str, str, Path]]: + copies: list[tuple[str, str, Path]] = [] + for relative in self.changed: + for form, path in _physical_paths(self.layout, self.root, relative).items(): + if path.is_file(): + copies.append((relative, form, path)) + return copies + + def test_every_d10_path_has_its_required_physical_form(self) -> None: + missing = [ + relative + for relative in self.changed + if not _canonical_paths(self.layout, self.root, relative).is_file() + ] + self.assertEqual(missing, [], f"{self.layout} canonical copies missing") + + def test_every_committed_d10_path_is_inside_the_d10_surface(self) -> None: + outside = [r for r in self.changed if not _in_d10_surface(r)] + self.assertEqual( + outside, + [], + "the committed D10 path list left the surface this gate scans; the " + "drift guard narrows source control to that surface, so a path " + "outside it would never be compared", + ) + + def test_committed_d10_paths_match_the_uncommitted_working_copy(self) -> None: + """Strict where source control can still identify D10; else says so.""" + reported = _surface_status_paths(self.layout, self.root) + if reported is None: + self.assertFalse( + _sapling_is_usable(self.root), + "`sl status` failed inside a Sapling working copy, so the " + "committed D10 path list went unverified here", + ) + return + if _THIS_GATE not in reported: + self.assertNotEqual( + _introducing_commit(self.root), + "", + "this gate is neither an uncommitted change nor a committed " + "file, so the committed D10 path list went unverified here", + ) + return + self.assertEqual( + sorted(reported), + list(self.changed), + "the committed D10 path list drifted from the working copy", + ) + + def test_fbcode_mirror_is_byte_identical_where_present(self) -> None: + if self.layout != "fbsource": + self.assertEqual(self.layout, "oss") + return + # Sapling mirrors xplat->fbcode on commit, so a path that is still an + # uncommitted working-copy change legitimately has no matching mirror. + uncommitted = _status_paths(self.root) + if uncommitted is None: + self.assertFalse( + _sapling_is_usable(self.root), + "`sl status` failed inside a Sapling working copy, so the paths " + "the mirror legitimately lags could not be identified", + ) + return + differing: list[str] = [] + absent_precommit: list[str] = [] + for relative in self.changed: + forms = _physical_paths(self.layout, self.root, relative) + xplat, fbcode = forms["xplat"], forms["fbcode"] + if not xplat.is_file(): + continue + if str(xplat.relative_to(self.root)) in uncommitted: + continue + if not fbcode.is_file(): + absent_precommit.append(relative) + continue + if xplat.read_bytes() != fbcode.read_bytes(): + differing.append(relative) + self.assertEqual( + differing, + [], + "xplat/fbcode mirror diverges; still-uncommitted paths are exempt " + f"and the mirror is absent for: {absent_precommit}", + ) + + def test_no_internal_only_reference_in_d10_sources(self) -> None: + violations: list[str] = [] + for relative, form, path in self._existing_copies(): + text = path.read_text(encoding="utf-8", errors="replace") + for number, line in enumerate(text.splitlines(), 1): + if _NEGATIVE_FIXTURE_MARKER in line: + continue + for pattern in _INTERNAL_REFERENCE_PATTERNS: + found = re.search(pattern, line, re.IGNORECASE) + if found is not None: + violations.append( + f"{form}:{relative}:{number}: {found.group(0)}" + ) + break + self.assertEqual(violations, [], "internal-only reference") + + def test_negative_fixture_marker_cannot_silence_other_files(self) -> None: + marked = { + relative + for relative, _form, path in self._existing_copies() + if _NEGATIVE_FIXTURE_MARKER + in path.read_text(encoding="utf-8", errors="replace") + } + self.assertEqual(marked, set(_NEGATIVE_FIXTURE_FILES)) + + def test_no_binary_artifact_in_d10_paths_or_directories(self) -> None: + untracked = _untracked_relatives(self.layout, self.root) + binaries = [ + relative + for relative in self.changed + if Path(relative).suffix in _BINARY_ARTIFACT_SUFFIXES + ] + for directory in _D10_SURFACE_DIRECTORIES: + for form, base in _physical_paths( + self.layout, self.root, directory + ).items(): + if not base.is_dir(): + continue + for path in sorted(base.rglob("*")): + if "__pycache__" in path.parts or not path.is_file(): + continue + if path.suffix not in _BINARY_ARTIFACT_SUFFIXES: + continue + if f"{directory}/{path.relative_to(base)}" in untracked: + continue + binaries.append(f"{form}:{path.relative_to(base)}") + self.assertEqual( + binaries, [], "committed binary artifact; untracked scratch is exempt" + ) + + def test_command_contract_module_paths_resolve(self) -> None: + unresolved: list[str] = [] + for module in _D10_COMMAND_MODULES: + self.assertTrue(module.startswith("executorch."), module) + relative = module[len("executorch.") :].replace(".", "/") + ".py" + if not _canonical_paths(self.layout, self.root, relative).is_file(): + unresolved.append(f"{module} -> {relative}") + for relative in _D10_COMMAND_SCRIPTS: + if not _canonical_paths(self.layout, self.root, relative).is_file(): + unresolved.append(relative) + self.assertEqual(unresolved, [], "command contract path does not resolve") + + def test_command_contract_unittest_target_resolves(self) -> None: + module, class_name, method = _D10_COMMAND_TEST_CASE + relative = module[len("executorch.") :].replace(".", "/") + ".py" + path = _canonical_paths(self.layout, self.root, relative) + self.assertTrue(path.is_file(), relative) + self.assertIsNotNone(_method_def(_module_tree(path), class_name, method)) + + +class PlainGemma4ContractTest(unittest.TestCase): + def setUp(self) -> None: + self.layout, self.root = _tree_root() + self.changed = _D10_CHANGED_PATHS + self.partitioner = _module_tree( + _canonical_paths(self.layout, self.root, _PLAIN_PARTITIONER) + ) + + def test_plain_partitioner_expected_counts_are_unchanged(self) -> None: + self.assertEqual( + _module_constant(self.partitioner, "_EXPECTED_GEMMA4_SDPA_COUNT"), 35 + ) + self.assertEqual( + _module_constant(self.partitioner, "_EXPECTED_SINGLE_HF_ROPE_COUNT"), 20 + ) + + def test_plain_webgpu_allowlist_membership_is_unchanged(self) -> None: + allowlist = _returned_dotted_names( + _function_def(self.partitioner, "_webgpu_allowlist") + ) + self.assertEqual(len(allowlist), len(_PLAIN_WEBGPU_ALLOWLIST)) + self.assertEqual(tuple(sorted(allowlist)), _PLAIN_WEBGPU_ALLOWLIST) + + def test_plain_webgpu_allowlist_has_no_mtp_operator(self) -> None: + allowlist = _returned_dotted_names( + _function_def(self.partitioner, "_webgpu_allowlist") + ) + forbidden = [ + entry + for entry in allowlist + if "scatter" in entry or "topk" in entry.lower() + ] + self.assertEqual(forbidden, [], "MTP operator leaked into plain allowlist") + + def test_plain_partitioner_fails_closed_on_emb8(self) -> None: + init = _method_def(self.partitioner, "Gemma4WebGPUPartitioner", "__init__") + guards = [ + statement + for statement in init.body + if isinstance(statement, ast.If) + and isinstance(statement.test, ast.Compare) + and isinstance(statement.test.left, ast.Constant) + and statement.test.left.value == "emb8" + ] + self.assertEqual(len(guards), 1) + guard = guards[0] + test = guard.test + self.assertIsInstance(test, ast.Compare) + assert isinstance(test, ast.Compare) + self.assertIsInstance(test.left, ast.Constant) + assert isinstance(test.left, ast.Constant) + self.assertEqual(test.left.value, "emb8") + self.assertEqual(len(test.ops), 1) + self.assertIsInstance(test.ops[0], ast.In) + self.assertEqual(_dotted_name(test.comparators[0]), "text_quantize") + raised = guard.body[0] + self.assertIsInstance(raised, ast.Raise) + assert isinstance(raised, ast.Raise) + self.assertIsInstance(raised.exc, ast.Call) + assert isinstance(raised.exc, ast.Call) + self.assertEqual(_dotted_name(raised.exc.func), "ValueError") + + def test_plain_manifest_identities_are_unchanged(self) -> None: + module = _load_plain_manifest_module(self.layout, self.root) + self.assertEqual(module.WEBGPU_EXTERNAL_CONSTANTS_MAX_DATA_BYTES, 1_500_000_000) + self.assertEqual(module.EXPORT_CONTRACT["methods"], ["text_decoder"]) + self.assertEqual(module.EXPORT_CONTRACT["max_input_len"], 512) + self.assertEqual(module.EXPORT_CONTRACT["max_seq_len"], 8960) + self.assertEqual(module.ARCHITECTURE_FINGERPRINT["num_hidden_layers"], 35) + self.assertEqual(len(module.ARCHITECTURE_FINGERPRINT["layer_types"]), 35) + self.assertEqual( + module.CHECKPOINT_ACQUISITION["repo_id"], + "google/gemma-4-E2B-it-qat-q4_0-unquantized", + ) + self.assertEqual( + _identity_sha256(module.EXPORT_CONTRACT), _PLAIN_EXPORT_CONTRACT_SHA256 + ) + self.assertEqual( + _identity_sha256(module.ARCHITECTURE_FINGERPRINT), + _PLAIN_ARCHITECTURE_SHA256, + ) + self.assertEqual( + _identity_sha256(module.CHECKPOINT_ACQUISITION), + _PLAIN_ACQUISITION_SHA256, + ) + + def test_plain_manifest_requires_exactly_three_ordered_ptds(self) -> None: + module = _load_plain_manifest_module(self.layout, self.root) + roles = {"pte": Path("model.pte"), "source": Path("source.json")} + for count in (0, 2, 4): + with self.assertRaises(ValueError) as raised: + module.create_plain_manifest( + self.root, roles, [Path(f"c{i}.ptd") for i in range(count)] + ) + self.assertIn("exactly three ordered PTDs", str(raised.exception)) + + def test_plain_manifest_json_matches_the_plain_validator(self) -> None: + module = _load_plain_manifest_module(self.layout, self.root) + path = _canonical_paths(self.layout, self.root, _PLAIN_MANIFEST_JSON) + manifest = json.loads(path.read_text(encoding="utf-8")) + + self.assertEqual(manifest["export"]["methods"], ["text_decoder"]) + self.assertEqual(manifest["export"], module.EXPORT_CONTRACT) + self.assertEqual(manifest["acquisition"], module.CHECKPOINT_ACQUISITION) + self.assertEqual( + manifest["model"]["architecture"], module.ARCHITECTURE_FINGERPRINT + ) + self.assertEqual( + manifest["model"]["source_config"], + { + "path": "config/e2b_config.json", + "sha256": module.SOURCE_CONFIG_SHA256, + }, + ) + self.assertEqual(manifest["schema_version"], 1) + + artifacts = manifest["artifacts"] + ptds = [item for item in artifacts if item["role"] == "ptd"] + self.assertEqual(len(manifest["ptd_order"]), 3) + self.assertEqual([item["path"] for item in ptds], manifest["ptd_order"]) + for item in ptds: + self.assertLess( + item["bytes"], module.WEBGPU_EXTERNAL_CONSTANTS_MAX_DATA_BYTES + ) + self.assertIn("pte", {item["role"] for item in artifacts}) + for item in artifacts: + self.assertEqual(len(Path(item["path"]).parts), 1, item["path"]) + + def test_public_xnnpack_runner_is_untouched_by_d10(self) -> None: + owned = [ + relative + for relative in self.changed + if relative.startswith("examples/models/gemma4/runner/") + ] + self.assertEqual(owned, [], "D10 must not own the public runner") + + header = _canonical_paths( + self.layout, self.root, _PLAIN_RUNNER_HEADER + ).read_text(encoding="utf-8") + self.assertIn("bool enable_workspace_sharing = true", header) + + source = _canonical_paths( + self.layout, self.root, _PLAIN_RUNNER_SOURCE + ).read_text(encoding="utf-8") + self.assertIn( + "#include ", + source, + ) + self.assertIn( + "executorch::backends::xnnpack::WorkspaceSharingMode::PerModel", + source, + ) + + targets = _canonical_paths( + self.layout, self.root, _PLAIN_MODEL_TARGETS + ).read_text(encoding="utf-8") + self.assertIn('"//executorch/backends/xnnpack:xnnpack_backend"', targets) + self.assertIn('"//executorch/backends/xnnpack:xnnpack_interface"', targets) diff --git a/examples/models/gemma4/tests/test_webgpu_artifact_manifest.py b/examples/models/gemma4/tests/test_webgpu_artifact_manifest.py index 68d05d9e604..1bb986b254b 100644 --- a/examples/models/gemma4/tests/test_webgpu_artifact_manifest.py +++ b/examples/models/gemma4/tests/test_webgpu_artifact_manifest.py @@ -6,97 +6,34 @@ import copy import hashlib +import inspect import json +import os +import re +import shutil +import signal import tempfile import unittest from pathlib import Path -from typing import Any from unittest import mock +from executorch.backends.webgpu.scripts import ( + webgpu_artifact_manifest as backend_manifest, +) from executorch.examples.models.gemma4 import ( + target_prefill_contract, webgpu_artifact_manifest as gemma4_manifest, ) - from executorch.examples.models.gemma4.webgpu_artifact_manifest import ( ARCHITECTURE_FINGERPRINT, + CHECKPOINT_ACQUISITION, create_plain_manifest, + EXPORT_CONTRACT, 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() @@ -124,9 +61,7 @@ def setUp(self) -> None: 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["ptd_order"], [path.name for path in self.ptds]) self.assertEqual( self.manifest["model"]["architecture"], ARCHITECTURE_FINGERPRINT ) @@ -224,6 +159,364 @@ def test_rejects_receipt_wgsl_head_mismatch(self) -> None: ) +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, object]: + logical_path = "examples/models/gemma4/webgpu_artifact_manifest.py" + identity = {"bytes": 7, "sha256": "3" * 64} + 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, + } + ] + return { + "checkouts": { + "fbsource": {"clean": True, "head": "1" * 40}, + "oss": {"clean": True, "head": "2" * 40}, + }, + "file_set_sha256": _set_digest([logical_path]), + "files": files, + "schema_version": 1, + } + + +def _test_wgsl_manifest() -> dict[str, object]: + 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, object]: + 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(), + } + + +_SEALED_SOURCE_RECEIPT = _sealed_source_receipt() +_BINARY_SUFFIXES = frozenset({".bin", ".gguf", ".pte", ".ptd", ".safetensors"}) +_ARTIFACT_KEYS = ["bytes", "path", "role", "sha256"] + + +def _write_source_receipt(path: Path) -> None: + path.write_text(json.dumps(_sealed_source_receipt()), encoding="utf-8") + + +def _identity(path: Path) -> dict[str, object]: + return { + "bytes": path.stat().st_size, + "sha256": hashlib.sha256(path.read_bytes()).hexdigest(), + } + + +def _target_prefill_receipt( + runtime_source_path: Path, +) -> dict[str, object]: + runtime_source = json.loads(runtime_source_path.read_text(encoding="utf-8")) + producer_path = target_prefill_contract.reviewed_producer_source_path() + contexts: dict[str, object] = {} + for context in target_prefill_contract.TARGET_PREFILL_CONTEXTS: + start, length = target_prefill_contract.final_chunk_range(context) + tensor = { + "byte_order": "little", + "dtype": "float32", + "layout": "row_major_contiguous", + "sha256": "a" * 64, + } + arm_config = { + "dtype": "float32", + "enable_dynamic_shape": True, + "group_size": 128, + "max_seq_len": 8960, + "text_quantize": "8da4w+emb4", + "use_kv_cache": True, + "variant": "e2b", + } + contexts[str(context)] = { + "arm_configs": { + "custom_sdpa_fused": {**arm_config, "use_custom_sdpa": True}, + "manual_unfused": {**arm_config, "use_custom_sdpa": 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, + "shape": [1, length, 8, 256], + }, + "manual_unfused": { + **tensor, + "sha256": "b" * 64, + "shape": [1, length, 8, 256], + }, + }, + "logits_post_softcap": { + **tensor, + "sha256": "c" * 64, + "shape": [1, 1, 262144], + }, + "logits_pre_softcap": { + **tensor, + "sha256": "d" * 64, + "shape": [1, 1, 262144], + }, + "prefill_token_post_softcap": 17, + "prefill_token_raw": 17, + "prompt_plan_sha256": target_prefill_contract.prompt_plan_sha256(context), + } + return { + "authority": "target_only_eager", + "checkpoint_acquisition": gemma4_manifest.CHECKPOINT_ACQUISITION, + "contexts": contexts, + "envelope_kind": "target_prefill_v2", + "producer": { + "fbsource_commit": runtime_source["fbsource_commit"], + "runtime_source_receipt": _identity(runtime_source_path), + "source_path": producer_path.name, + "source_sha256": hashlib.sha256(producer_path.read_bytes()).hexdigest(), + }, + "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, + } + + +def _generated_mtp_evidence() -> dict[str, object]: + mutation_order: list[dict[str, object]] = [ + { + "logicalTarget": "seed_feature", + "role": "nextFeatureSeed", + "shape": [1, 1, 1, 1536], + "logicalLayout": "BSHD", + "logicalDimOrder": [0, 1, 2, 3], + "vulkanSourceStorage": "BUFFER", + "vulkanDestinationStorage": "TEXTURE_3D", + } + ] + for layer in range(15): + head_dim = 512 if layer in {4, 9, 14} else 256 + for cache_kind in ("k_cache", "v_cache"): + mutation_order.append( + { + "logicalTarget": ( + f"self_decoder.layers.{layer}.self_attn.kv_cache." + f"{cache_kind}" + ), + "role": "targetKvCache", + "layer": layer, + "cacheKind": cache_kind, + "shape": [1, 8960, 1, head_dim], + "logicalLayout": "BSHD", + "logicalDimOrder": [0, 1, 2, 3], + "vulkanSourceStorage": "BUFFER", + "vulkanDestinationStorage": "BUFFER", + } + ) + token_record: dict[str, object] = { + "max": 262143, + "min": 0, + "numel": 262144, + "permutationExact": True, + "rawShape": [262144], + "sha256": "5" * 64, + "shape": [2048, 128], + "uniqueCount": 262144, + } + token_ordering = { + **token_record, + "loaded": copy.deepcopy(token_record), + "raw": copy.deepcopy(token_record), + "rawLoadedByteExact": True, + "rawSha256": token_record["sha256"], + } + donor_sequence = [2, 16, 511, 512, 513, 514, 1024, 8960, 2] + cases = [ + { + "caseIndex": index, + "donorLength": donor_length, + "greedyTokenExact": True, + "inputSha256": ["6" * 64], + "outputs": [ + { + "actualSha256": "7" * 64, + "bitExact": True, + "close": True, + "maxAbsError": 0.0, + "name": name, + "referenceSha256": "7" * 64, + "shape": [1, 1], + } + for name in ("logits", "last_hidden_state") + ], + "topk": { + "allFinite": True, + "boundaryGap": 1.0, + "indicesSha256": "8" * 64, + "stableReferenceExact": True, + "top32PairwiseDistinct": True, + "top33IndicesSha256": "9" * 64, + "top33ValuesSha256": "a" * 64, + "valuesSha256": "b" * 64, + }, + } + for index, donor_length in enumerate(donor_sequence) + ] + return { + "assistant_checkpoint": gemma4_manifest.ASSISTANT_CHECKPOINT_ACQUISITION, + "k2_abi": { + "bufferMutationCount": 31, + "donorViewOrder": [ + { + "role": "fullK", + "layer": 14, + "cacheKind": "k_cache", + "layout": "BHKD", + }, + { + "role": "fullV", + "layer": 14, + "cacheKind": "v_cache", + "layout": "BHKD", + }, + { + "role": "slidingK", + "layer": 13, + "cacheKind": "k_cache", + "layout": "BHKD", + }, + { + "role": "slidingV", + "layer": 13, + "cacheKind": "v_cache", + "layout": "BHKD", + }, + ], + "inputOrder": ["input_ids", "input_pos", "is_round", "donor_length"], + "mutationOrder": mutation_order, + "operatorCounts": { + "aten.argmax.default": 3, + "aten.scatter.src": 2, + "aten.topk.default": 2, + "llama.custom_sdpa.default": 43, + "llama.update_cache.default": 31, + }, + "outputOrder": [ + "candidates", + "target_greedy", + "output_matches", + "output_bonus", + "state_probe", + ], + "seedMutationCount": 1, + "stateAlias": { + "logicalSource": "nextFeature[1,1,1536]", + "physicalDestination": "seed_feature[1,1,1,1536]", + "mutation": "llama.update_cache.default", + }, + }, + "lowering": { + "delegate_count": 1, + "edge": gemma4_manifest.MTP_EDGE_CENSUS, + "portable_operator_count": 0, + }, + "qat_selection": { + "cases": cases, + "donorSequence": donor_sequence, + "eagerEquivalence": {"allClose": True, "atol": 1e-4, "rtol": 1e-3}, + "selectionContract": { + "centroidTopK": 32, + "numCentroids": 2048, + "selectedTokenCount": 4096, + "tokensPerCentroid": 128, + }, + "tokenOrdering": token_ordering, + }, + "target_checkpoint": gemma4_manifest.CHECKPOINT_ACQUISITION, + } + + +def _artifact_entry(manifest: dict[str, object], path: str) -> dict[str, object]: + artifacts = manifest["artifacts"] + matches = [entry for entry in artifacts if entry["path"] == path] + if len(matches) != 1: + raise AssertionError(f"expected exactly one artifact for {path}") + return matches[0] + + +def _package_root() -> Path: + return Path(gemma4_manifest.__file__).resolve().parent + + +def _internal_patterns() -> list[re.Pattern[str]]: + return [ + re.compile(pattern) + for pattern in ( + r"manifold", # oss-closure-fixture + r"internalfb", # oss-closure-fixture + r"/data/users/", # oss-closure-fixture + r"fburl", # oss-closure-fixture + r"/home/", # oss-closure-fixture + r"\bfb/", + ) + ] + + class SourceClosureManifestTest(unittest.TestCase): def setUp(self) -> None: self._temporary_directory = tempfile.TemporaryDirectory() @@ -233,16 +526,21 @@ def setUp(self) -> None: 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") + self.owned_paths = [ + "examples/models/gemma4/webgpu_artifact_manifest.py", + "backends/webgpu/runtime/WebGPUShaderRegistry.cpp", + ] + for index, logical_path in enumerate(self.owned_paths): + contents = f"source-{index}".encode("utf-8") + for path in ( + self.fbsource_root / "fbcode/executorch" / logical_path, + self.fbsource_root / "xplat/executorch" / logical_path, + self.oss_root / logical_path, + ): + path.parent.mkdir(parents=True, exist_ok=True) + path.write_bytes(contents) - def _create_source_manifest(self) -> dict[str, Any]: + def _create_source_manifest(self) -> dict[str, object]: def snapshot(_root: Path, kind: str) -> dict[str, object]: return { "clean": True, @@ -254,53 +552,13 @@ def snapshot(_root: Path, kind: str) -> dict[str, object]: ), mock.patch.object( gemma4_manifest, "_derive_owned_paths", - return_value=[self.logical_path], + return_value=sorted(self.owned_paths), ): 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() @@ -321,8 +579,9 @@ def test_create_source_manifest_cli_round_trip(self) -> None: ), 0, ) - document = json.loads(output.read_text(encoding="utf-8")) - gemma4_manifest.validate_source_manifest(document) + gemma4_manifest.validate_source_manifest( + json.loads(output.read_text(encoding="utf-8")) + ) def test_create_wgsl_manifest_cli_round_trip(self) -> None: output = self.root / "wgsl.json" @@ -342,14 +601,17 @@ def test_create_wgsl_manifest_cli_round_trip(self) -> None: ), 0, ) - document = json.loads(output.read_text(encoding="utf-8")) - gemma4_manifest.validate_wgsl_manifest(document) + gemma4_manifest.validate_wgsl_manifest( + json.loads(output.read_text(encoding="utf-8")) + ) 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 + gemma4_manifest, + "create_source_closure_receipt", + return_value=receipt, ): self.assertEqual( gemma4_manifest.main( @@ -369,37 +631,2357 @@ def test_create_source_receipt_cli_round_trip(self) -> None: ) 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"), + def _copy_backend_root(self) -> Path: + source = Path(backend_manifest.__file__).resolve().parents[1] + backend_root = ( + self.root / "xplat" / "executorch" / "backends" / "webgpu" + ) + shutil.copytree(source, backend_root) + return backend_root + + def test_source_manifest_does_not_accept_a_caller_selected_subset(self) -> None: + parameters = inspect.signature( + gemma4_manifest.create_source_manifest + ).parameters + self.assertNotIn("owned_paths", parameters) + + def test_owned_union_is_derived_from_every_reviewed_diff(self) -> None: + production_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", + "[ExecuTorch][WebGPU] Add Gemma 4 MTP operator and route support", + "[ExecuTorch][WebGPU] Add Gemma 4 MTP export path", + "[ExecuTorch][WebGPU] Add Gemma 4 speculative decode runtime", + "[ExecuTorch][WebGPU] Add Gemma 4 MTP and speculative-decode source-closure tests", + ) + self.assertEqual( + gemma4_manifest._GEMMA_PRODUCTION_DIFF_SUMMARIES, + production_summaries, + ) + summaries = tuple(reversed(production_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{summaries[offset]}\n" + node = argv[argv.index("--change") + 1] + offset = int(node, 16) - 1 + path = f"runtime/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 ): - path.parent.mkdir(parents=True, exist_ok=True) - path.write_bytes(contents) + paths = gemma4_manifest._derive_owned_paths(self.fbsource_root) + self.assertEqual( + paths, + [f"runtime/owned_{index}.cpp" for index in range(9)], + ) + + def test_source_manifest_derives_heads_and_binds_every_copy(self) -> None: + manifest = self._create_source_manifest() + gemma4_manifest.validate_source_manifest(manifest) + self.assertEqual( + manifest["checkouts"], + { + "fbsource": {"clean": True, "head": "1" * 40}, + "oss": {"clean": True, "head": "2" * 40}, + }, + ) + files = manifest["files"] + self.assertEqual([entry["path"] for entry in files], sorted(self.owned_paths)) + for entry in files: + identities = { + (copy["bytes"], copy["sha256"]) for copy in entry["copies"].values() + } + self.assertEqual(len(identities), 1) + + def test_source_manifest_rejects_mirror_or_oss_drift(self) -> None: + (self.oss_root / self.owned_paths[0]).write_bytes(b"different") + with self.assertRaisesRegex(ValueError, "mirror/OSS identity mismatch"): + self._create_source_manifest() + + def test_source_manifest_rejects_ancestor_symlink_traversal(self) -> None: + parent = self.fbsource_root / "xplat/executorch/examples" + saved = self.root / "saved-examples" + parent.rename(saved) + parent.symlink_to(saved, target_is_directory=True) + with self.assertRaisesRegex(ValueError, "symlink traversal"): + self._create_source_manifest() - class Generator: - def discover(self) -> list[Path]: - return [shader] + def test_source_manifest_fails_closed_without_clean_oss_identity(self) -> None: + def snapshot(_root: Path, kind: str) -> dict[str, object]: + if kind == "oss": + raise ValueError("oss checkout is not clean") + return {"clean": True, "head": "1" * 40} - def collect_outputs(self) -> tuple[dict[Path, bytes], list[Path]]: - return {header: b"fresh", registry: b"registry"}, [] + with mock.patch.object( + gemma4_manifest, "_checkout_snapshot", side_effect=snapshot + ): + with self.assertRaisesRegex(ValueError, "oss checkout is not clean"): + gemma4_manifest.create_source_manifest( + self.fbsource_root, self.oss_root + ) - def registry_path(self) -> Path: - return registry + def test_wgsl_manifest_uses_generator_complete_dynamic_closure(self) -> None: + backend_root = self._copy_backend_root() + with mock.patch.object( + gemma4_manifest, + "_checkout_snapshot", + return_value={"clean": True, "head": "1" * 40}, + ): + manifest = gemma4_manifest.create_wgsl_manifest(backend_root) + gemma4_manifest.validate_wgsl_manifest(manifest) + self.assertEqual(manifest["fbsource_commit"], "1" * 40) + roles = [entry["role"] for entry in manifest["files"]] + self.assertEqual(roles.count("generator"), 1) + self.assertEqual(roles.count("global_registry"), 1) + self.assertGreater(roles.count("wgsl"), 0) + self.assertGreater(roles.count("generated_header"), roles.count("wgsl")) + self.assertEqual(manifest["orphans"], []) + def test_wgsl_manifest_rejects_generator_output_byte_drift(self) -> None: + backend_root = self._copy_backend_root() + generator = gemma4_manifest._load_wgsl_generator(backend_root) + outputs, orphans = generator.collect_outputs() + stale_outputs = dict(outputs) + path = next(iter(stale_outputs)) + generated = stale_outputs[path] + stale_outputs[path] = bytes([generated[0] ^ 1]) + generated[1:] 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() + gemma4_manifest, "_load_wgsl_generator", return_value=generator + ), mock.patch.object( + generator, "collect_outputs", return_value=(stale_outputs, orphans) ): with self.assertRaisesRegex(ValueError, "generated output is stale"): gemma4_manifest.create_wgsl_manifest(backend_root) + + def test_source_receipt_is_derived_from_live_roots_not_manifest_json(self) -> None: + parameters = inspect.signature( + gemma4_manifest.create_source_closure_receipt + ).parameters + self.assertNotIn("source_manifest_path", parameters) + self.assertNotIn("wgsl_manifest_path", parameters) + with mock.patch.object( + gemma4_manifest, + "create_source_manifest", + return_value=_test_source_manifest(), + ) as create_source, mock.patch.object( + gemma4_manifest, + "create_wgsl_manifest", + return_value=_test_wgsl_manifest(), + ) as create_wgsl: + receipt = gemma4_manifest.create_source_closure_receipt( + self.fbsource_root, self.oss_root, self.fbsource_root + ) + create_source.assert_called_once_with(self.fbsource_root, self.oss_root) + create_wgsl.assert_called_once_with(self.fbsource_root) + self.assertEqual(receipt["fbsource_commit"], "1" * 40) + + def test_source_receipt_rejects_wgsl_from_another_checkout(self) -> None: + wgsl = _test_wgsl_manifest() + wgsl["fbsource_commit"] = "9" * 40 + with mock.patch.object( + gemma4_manifest, + "create_source_manifest", + return_value=_test_source_manifest(), + ), mock.patch.object( + gemma4_manifest, "create_wgsl_manifest", return_value=wgsl + ): + with self.assertRaisesRegex(ValueError, "different fbsource heads"): + gemma4_manifest.create_source_closure_receipt( + self.fbsource_root, self.oss_root, self.fbsource_root + ) + + +class MTPArtifactManifestTest(unittest.TestCase): + def setUp(self) -> None: + # Deferred: the plain cases above must load without the D8 manifest extension. + from executorch.examples.models.gemma4.webgpu_artifact_manifest import ( + create_mtp_manifest, + validate_mtp_manifest, + ) + + self.create_mtp_manifest = create_mtp_manifest + self.validate_mtp_manifest = validate_mtp_manifest + self._temporary_directory = tempfile.TemporaryDirectory() + self.addCleanup(self._temporary_directory.cleanup) + self.root = Path(self._temporary_directory.name) + self.pte = self.root / "k2_round.pte" + self.ptds = [self.root / f"constants_{index}.ptd" for index in range(3)] + # Written only by the source-verified cases: staging must stay exact. + self.source_receipt = self.root / "source_receipt.json" + self.pte.write_bytes(b"k2-round-pte") + for index, path in enumerate(self.ptds): + path.write_bytes(f"mtp-ptd-{index}".encode("utf-8")) + self.manifest = self.create_mtp_manifest( + self.root, + {"pte": Path(self.pte.name)}, + [Path(path.name) for path in self.ptds], + ) + self.manifest["provenance"] = copy.deepcopy( + gemma4_manifest.MTP_ACCEPTED_PROVENANCE + ) + + def _pending_manifest(self) -> dict[str, object]: + manifest = self.create_mtp_manifest( + self.root, + {"pte": Path(self.pte.name)}, + [Path(path.name) for path in self.ptds], + ) + manifest["evidence"] = _generated_mtp_evidence() + return manifest + + def _verified_manifest(self) -> dict[str, object]: + _write_source_receipt(self.source_receipt) + manifest = self.create_mtp_manifest( + self.root, + { + "pte": Path(self.pte.name), + "source": Path(self.source_receipt.name), + }, + [Path(path.name) for path in self.ptds], + ) + manifest["evidence"] = _generated_mtp_evidence() + return manifest + + def test_single_k2_pte_and_three_ptds_round_trip(self) -> None: + self.validate_mtp_manifest(self.root, self.manifest) + artifacts = self.manifest["artifacts"] + self.assertEqual( + sorted(entry["role"] for entry in artifacts), + ["ptd", "ptd", "ptd", "pte"], + ) + for entry in artifacts: + self.assertEqual(sorted(entry), _ARTIFACT_KEYS) + self.assertRegex(str(entry["sha256"]), "^[0-9a-f]{64}$") + self.assertEqual( + _artifact_entry(self.manifest, "k2_round.pte")["sha256"], + hashlib.sha256(b"k2-round-pte").hexdigest(), + ) + self.assertEqual(_artifact_entry(self.manifest, "k2_round.pte")["bytes"], 12) + + def test_create_mtp_cli_requires_evidence_and_emits_a_valid_manifest( + self, + ) -> None: + arguments = [ + "create-mtp", + "--root", + str(self.root), + "--role", + f"pte={self.pte.name}", + *(argument for path in self.ptds for argument in ("--ptd", path.name)), + ] + with mock.patch("sys.stderr"), self.assertRaises(SystemExit): + gemma4_manifest.main([*arguments, "--output", str(self.root / "unused")]) + + with tempfile.TemporaryDirectory() as metadata_directory: + metadata_root = Path(metadata_directory) + evidence_path = metadata_root / "evidence.json" + output_path = metadata_root / "mtp.json" + evidence = _generated_mtp_evidence() + evidence_path.write_text(json.dumps(evidence), encoding="utf-8") + self.assertEqual( + gemma4_manifest.main( + [ + *arguments, + "--evidence", + str(evidence_path), + "--output", + str(output_path), + ] + ), + 0, + ) + manifest = json.loads(output_path.read_text(encoding="utf-8")) + self.assertEqual(manifest["evidence"], evidence) + self.validate_mtp_manifest(self.root, manifest) + + def test_mtp_creation_rejects_split_model_roles(self) -> None: + for role in ("assistant", "speculative"): + with self.subTest(role=role): + with self.assertRaisesRegex(ValueError, "one K=2 PTE role"): + self.create_mtp_manifest( + self.root, + {"pte": Path(self.pte.name), role: Path(self.pte.name)}, + [Path(path.name) for path in self.ptds], + ) + + def test_mtp_creation_rejects_duplicate_normalized_artifact_paths(self) -> None: + ptd_paths = [Path(path.name) for path in self.ptds] + for source_alias in (self.pte, *self.ptds): + with self.subTest(source_alias=source_alias.name): + with self.assertRaisesRegex( + ValueError, "duplicate normalized artifact path" + ): + self.create_mtp_manifest( + self.root, + { + "pte": Path(self.pte.name), + "source": Path(source_alias.name), + }, + ptd_paths, + ) + + def test_mtp_validation_rejects_duplicate_normalized_artifact_paths( + self, + ) -> None: + for source_alias in (self.pte, *self.ptds): + with self.subTest(source_alias=source_alias.name): + duplicated = self._verified_manifest() + alias_entry = _artifact_entry(duplicated, source_alias.name) + source_entry = _artifact_entry(duplicated, self.source_receipt.name) + source_entry.update( + { + "bytes": alias_entry["bytes"], + "path": alias_entry["path"], + "sha256": alias_entry["sha256"], + } + ) + with self.assertRaisesRegex( + ValueError, "duplicate normalized artifact path" + ): + self.validate_mtp_manifest(self.root, duplicated) + + def test_mtp_ptd_order_is_the_artifact_order(self) -> None: + artifacts = self.manifest["artifacts"] + self.assertEqual( + self.manifest["ptd_order"], + [entry["path"] for entry in artifacts if entry["role"] == "ptd"], + ) + self.assertEqual(self.manifest["ptd_order"], [path.name for path in self.ptds]) + + reordered = copy.deepcopy(self.manifest) + order = reordered["ptd_order"] + order[0], order[1] = order[1], order[0] + with self.assertRaisesRegex(ValueError, "PTD order does not match"): + self.validate_mtp_manifest(self.root, reordered) + + def test_mtp_rejects_symlinked_artifact(self) -> None: + self.pte.unlink() + self.pte.symlink_to(self.ptds[0].name) + with self.assertRaisesRegex(ValueError, "symlink"): + self.validate_mtp_manifest(self.root, self.manifest) + + def test_mtp_rejects_absolute_artifact_path(self) -> None: + mutated = copy.deepcopy(self.manifest) + _artifact_entry(mutated, "k2_round.pte")["path"] = str(self.pte) + with self.assertRaisesRegex(ValueError, "absolute artifact path"): + self.validate_mtp_manifest(self.root, mutated) + + def test_mtp_rejects_non_canonical_artifact_path(self) -> None: + mutated = copy.deepcopy(self.manifest) + _artifact_entry(mutated, "k2_round.pte")["path"] = "./k2_round.pte" + with self.assertRaisesRegex(ValueError, "non-canonical artifact path"): + self.validate_mtp_manifest(self.root, mutated) + + def test_mtp_rejects_wrong_size(self) -> None: + self.pte.write_bytes(b"k2-round-pte-grew") + with self.assertRaisesRegex(ValueError, "byte count mismatch"): + self.validate_mtp_manifest(self.root, self.manifest) + + def test_mtp_rejects_wrong_hash(self) -> None: + self.pte.write_bytes(b"k2-round-ptX") + with self.assertRaisesRegex(ValueError, "SHA-256 mismatch"): + self.validate_mtp_manifest(self.root, self.manifest) + + def test_mtp_rejects_missing_artifact(self) -> None: + self.ptds[2].unlink() + with self.assertRaises(FileNotFoundError): + self.validate_mtp_manifest(self.root, self.manifest) + + def test_mtp_rejects_extra_artifact(self) -> None: + extra = self.root / "constants_3.ptd" + extra.write_bytes(b"mtp-ptd-3") + mutated = copy.deepcopy(self.manifest) + mutated["artifacts"].append( + { + "bytes": extra.stat().st_size, + "path": extra.name, + "role": "ptd", + "sha256": hashlib.sha256(extra.read_bytes()).hexdigest(), + } + ) + with self.assertRaisesRegex(ValueError, "PTD order does not match"): + self.validate_mtp_manifest(self.root, mutated) + + def test_mtp_rejects_stale_staging(self) -> None: + (self.root / "stale_constants.ptd").write_bytes(b"stale") + with self.assertRaisesRegex(ValueError, "missing or extra"): + self.validate_mtp_manifest(self.root, self.manifest) + + def test_pending_source_manifest_is_a_valid_state(self) -> None: + pending = self._pending_manifest() + self.assertEqual( + pending["provenance"], gemma4_manifest.MTP_PENDING_SOURCE_PROVENANCE + ) + self.validate_mtp_manifest(self.root, pending) + + def test_pending_source_manifest_is_not_source_complete(self) -> None: + with self.assertRaisesRegex(ValueError, "source closure is still pending"): + gemma4_manifest._validate_source_complete_mtp_manifest( + self.root, self._pending_manifest() + ) + + def test_sealed_source_receipt_stamps_source_verified(self) -> None: + verified = self._verified_manifest() + self.assertEqual( + verified["provenance"], gemma4_manifest.MTP_SOURCE_VERIFIED_PROVENANCE + ) + self.validate_mtp_manifest(self.root, verified) + gemma4_manifest._validate_source_complete_mtp_manifest(self.root, verified) + + def test_tampered_source_receipt_bytes_are_rejected(self) -> None: + verified = self._verified_manifest() + self.source_receipt.write_text( + json.dumps({**_SEALED_SOURCE_RECEIPT, "oss_commit": "3" * 40}), + encoding="utf-8", + ) + with self.assertRaisesRegex(ValueError, "SHA-256 mismatch"): + self.validate_mtp_manifest(self.root, verified) + + def test_unsealed_source_receipt_is_rejected_after_rehashing(self) -> None: + verified = self._verified_manifest() + self.source_receipt.write_text( + json.dumps({**_SEALED_SOURCE_RECEIPT, "source_current": False}), + encoding="utf-8", + ) + entry = _artifact_entry(verified, self.source_receipt.name) + entry["bytes"] = self.source_receipt.stat().st_size + entry["sha256"] = hashlib.sha256(self.source_receipt.read_bytes()).hexdigest() + with self.assertRaisesRegex(ValueError, "not source-current"): + self.validate_mtp_manifest(self.root, verified) + + def test_source_receipt_field_checks_are_enforced(self) -> None: + mutations = ( + (("schema_version",), 2, "not source-current"), + (("fbsource_commit",), "g" * 40, "invalid fbsource commit"), + (("oss_commit",), "z" * 40, "invalid OSS commit"), + ( + ("source_manifest", "file_set_sha256"), + "X" * 64, + "source manifest file-set identity", + ), + ( + ("wgsl_manifest", "file_set_sha256"), + "q" * 64, + "WGSL manifest file-set identity", + ), + ( + ("verification", "source_checkout"), + "pending", + "verification is incomplete", + ), + ) + for keys, value, message in mutations: + with self.subTest(keys=keys): + receipt = copy.deepcopy(_SEALED_SOURCE_RECEIPT) + target = receipt + for key in keys[:-1]: + target = target[key] + target[keys[-1]] = value + self.source_receipt.write_text(json.dumps(receipt), encoding="utf-8") + with self.assertRaisesRegex(ValueError, message): + self.create_mtp_manifest( + self.root, + { + "pte": Path(self.pte.name), + "source": Path(self.source_receipt.name), + }, + [Path(path.name) for path in self.ptds], + ) + + def test_mtp_rejects_fake_semantic_evidence(self) -> None: + mutations = { + "negative mutation count": lambda evidence: evidence["k2_abi"].update( + {"bufferMutationCount": -1} + ), + "empty input order": lambda evidence: evidence["k2_abi"].update( + {"inputOrder": []} + ), + "false state alias": lambda evidence: evidence["k2_abi"].update( + {"stateAlias": False} + ), + "negative operator count": lambda evidence: evidence["k2_abi"][ + "operatorCounts" + ].update({"aten.argmax.default": -1}), + "empty QAT cases": lambda evidence: evidence["qat_selection"].update( + {"cases": []} + ), + "invalid token digest": lambda evidence: evidence["qat_selection"][ + "tokenOrdering" + ].update({"sha256": "not-a-digest", "rawSha256": "not-a-digest"}), + "incoherent token statistics": lambda evidence: evidence["qat_selection"][ + "tokenOrdering" + ].update({"min": 1, "max": 2, "numel": 3, "uniqueCount": 4}), + } + for label, mutate in mutations.items(): + with self.subTest(label=label): + manifest = self._verified_manifest() + evidence = manifest["evidence"] + assert isinstance(evidence, dict) + mutate(evidence) + with self.assertRaises(ValueError): + self.validate_mtp_manifest(self.root, manifest) + + def test_forced_verified_label_without_a_receipt_is_rejected(self) -> None: + forced = self._pending_manifest() + forced["provenance"] = copy.deepcopy( + gemma4_manifest.MTP_SOURCE_VERIFIED_PROVENANCE + ) + with self.assertRaisesRegex(ValueError, "unexpected artifact roles"): + self.validate_mtp_manifest(self.root, forced) + + def test_forced_pending_label_with_a_receipt_is_rejected(self) -> None: + forced = self._verified_manifest() + forced["provenance"] = copy.deepcopy( + gemma4_manifest.MTP_PENDING_SOURCE_PROVENANCE + ) + with self.assertRaisesRegex(ValueError, "unexpected artifact roles"): + self.validate_mtp_manifest(self.root, forced) + + def test_source_verified_is_distinct_from_the_pending_closures(self) -> None: + verified = gemma4_manifest.MTP_SOURCE_VERIFIED_PROVENANCE["source_closure"] + pending = gemma4_manifest.MTP_PENDING_SOURCE_PROVENANCE["source_closure"] + self.assertEqual(verified, "source_verified") + self.assertEqual(pending, "pending_final_source_receipt") + self.assertNotEqual(verified, pending) + self.assertNotEqual( + gemma4_manifest.MTP_SOURCE_VERIFIED_PROVENANCE, + gemma4_manifest.MTP_PENDING_SOURCE_PROVENANCE, + ) + self.assertIn(pending, gemma4_manifest.MTP_PENDING_SOURCE_CLOSURES) + self.assertIn( + gemma4_manifest.MTP_ACCEPTED_PROVENANCE["source_closure"], + gemma4_manifest.MTP_PENDING_SOURCE_CLOSURES, + ) + self.assertNotIn(verified, gemma4_manifest.MTP_PENDING_SOURCE_CLOSURES) + + +class MTPExportPublicationTest(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.artifact_root = self.root / "sealed" + self.output_path = self.artifact_root / "model.pte" + self.receipt_path = self.root / "receipts" / "manifest.json" + self.source_root = self.root / "source" + self.source_root.mkdir() + self.source_receipt = self.source_root / "source_receipt.json" + _write_source_receipt(self.source_receipt) + + def _write_staged_artifacts(self, staging: Path) -> tuple[Path, list[Path]]: + staged_pte = staging / self.output_path.name + staged_ptds = [staging / f"constants_{index}.ptd" for index in range(3)] + staged_pte.write_bytes(b"mtp-pte") + for index, path in enumerate(staged_ptds): + path.write_bytes(f"mtp-ptd-{index}".encode("utf-8")) + return staged_pte, staged_ptds + + def _finalize( + self, + staging: Path, + staged_pte: Path, + staged_ptds: list[Path], + source_receipt: Path | None = None, + ) -> Path: + self.assertTrue(hasattr(gemma4_manifest, "finalize_mtp_export")) + return gemma4_manifest.finalize_mtp_export( + staging, + self.output_path, + self.receipt_path, + staged_pte, + staged_ptds, + source_receipt or self.source_receipt, + _generated_mtp_evidence(), + ) + + def _published_paths(self) -> tuple[Path, ...]: + return ( + *(self.artifact_root / f"constants_{index}.ptd" for index in range(3)), + self.artifact_root / self.source_receipt.name, + self.output_path, + self.receipt_path, + ) + + def _assert_no_publications(self) -> None: + for path in self._published_paths(): + with self.subTest(unpublished=path.name): + self.assertFalse(path.exists() or path.is_symlink()) + + def _assert_validation_failure_rolls_back(self, failure_call: int) -> None: + real_validate = gemma4_manifest.validate_mtp_manifest + call_count = 0 + + def validate_then_fail(root: Path, manifest: dict[str, object]) -> None: + nonlocal call_count + call_count += 1 + real_validate(root, manifest) + if call_count == failure_call: + self.assertEqual(self.receipt_path.is_file(), failure_call == 3) + raise ValueError(f"injected validation failure at call {failure_call}") + + with tempfile.TemporaryDirectory(dir=self.root) as directory: + staging = Path(directory) + staged_pte, staged_ptds = self._write_staged_artifacts(staging) + with mock.patch.object( + gemma4_manifest, + "validate_mtp_manifest", + side_effect=validate_then_fail, + ) as validator: + with self.assertRaisesRegex( + ValueError, + f"injected validation failure at call {failure_call}", + ): + self._finalize(staging, staged_pte, staged_ptds) + + self.assertEqual(validator.call_count, failure_call) + self._assert_no_publications() + + def test_source_receipt_survives_staging_cleanup_and_final_manifest_validates( + self, + ) -> None: + with tempfile.TemporaryDirectory(dir=self.root) as directory: + staging = Path(directory) + staged_pte, staged_ptds = self._write_staged_artifacts(staging) + result = self._finalize(staging, staged_pte, staged_ptds) + + self.assertEqual(result, self.receipt_path) + self.assertFalse(staging.exists()) + published_source = self.artifact_root / self.source_receipt.name + self.assertEqual( + published_source.read_bytes(), self.source_receipt.read_bytes() + ) + receipt = json.loads(self.receipt_path.read_text(encoding="utf-8")) + gemma4_manifest.validate_mtp_manifest(self.artifact_root, receipt) + self.assertEqual( + sorted(path.name for path in self.artifact_root.iterdir()), + [ + "constants_0.ptd", + "constants_1.ptd", + "constants_2.ptd", + "model.pte", + "source_receipt.json", + ], + ) + + def test_final_validation_failure_rolls_back_every_published_artifact( + self, + ) -> None: + self.artifact_root.mkdir() + extra = self.artifact_root / "unsealed.bin" + extra.write_bytes(b"not part of the sealed export") + with tempfile.TemporaryDirectory(dir=self.root) as directory: + staging = Path(directory) + staged_pte, staged_ptds = self._write_staged_artifacts(staging) + with self.assertRaisesRegex(ValueError, "missing or extra"): + self._finalize(staging, staged_pte, staged_ptds) + + self._assert_no_publications() + self.assertEqual(list(self.artifact_root.iterdir()), [extra]) + + def test_second_validation_failure_rolls_back_every_published_artifact( + self, + ) -> None: + self._assert_validation_failure_rolls_back(2) + + def test_third_validation_failure_rolls_back_published_receipt_and_artifacts( + self, + ) -> None: + self._assert_validation_failure_rolls_back(3) + + def test_keyboard_interrupt_rolls_back_every_published_artifact(self) -> None: + real_validate = gemma4_manifest.validate_mtp_manifest + call_count = 0 + + def validate_then_interrupt(root: Path, manifest: dict[str, object]) -> None: + nonlocal call_count + call_count += 1 + real_validate(root, manifest) + if call_count == 2: + raise KeyboardInterrupt("injected publication interrupt") + + with tempfile.TemporaryDirectory(dir=self.root) as directory: + staging = Path(directory) + staged_pte, staged_ptds = self._write_staged_artifacts(staging) + with mock.patch.object( + gemma4_manifest, + "validate_mtp_manifest", + side_effect=validate_then_interrupt, + ): + with self.assertRaisesRegex( + KeyboardInterrupt, "injected publication interrupt" + ): + self._finalize(staging, staged_pte, staged_ptds) + self._assert_no_publications() + + def _assert_post_link_sigint_rolls_back( + self, interrupted_destination: Path + ) -> None: + real_link = os.link + + def link_then_sigint(source: Path, target: Path, **kwargs: object) -> None: + real_link(source, target, **kwargs) + if target == interrupted_destination: + os.kill(os.getpid(), signal.SIGINT) + + with tempfile.TemporaryDirectory(dir=self.root) as directory: + staging = Path(directory) + staged_pte, staged_ptds = self._write_staged_artifacts(staging) + previous_handler = signal.signal(signal.SIGINT, signal.default_int_handler) + try: + with mock.patch.object( + os, + "link", + side_effect=link_then_sigint, + ): + with self.assertRaises(KeyboardInterrupt): + self._finalize(staging, staged_pte, staged_ptds) + finally: + signal.signal(signal.SIGINT, previous_handler) + + self._assert_no_publications() + self.assertTrue(self.source_receipt.is_file()) + + def test_post_link_sigint_is_deferred_and_rolls_back_each_destination(self) -> None: + for interrupted_destination in self._published_paths(): + with self.subTest(destination=interrupted_destination.name): + self._assert_post_link_sigint_rolls_back(interrupted_destination) + + def test_post_link_sigint_fixture_restores_inherited_handler(self) -> None: + inherited_signals: list[int] = [] + + def inherited_handler(signum: int, _frame: object) -> None: + inherited_signals.append(signum) + + previous_handler = signal.signal(signal.SIGINT, inherited_handler) + try: + self._assert_post_link_sigint_rolls_back(self.output_path) + self.assertIs(signal.getsignal(signal.SIGINT), inherited_handler) + finally: + signal.signal(signal.SIGINT, previous_handler) + + self.assertEqual(inherited_signals, []) + + def test_post_link_sigint_redelivers_to_inherited_handler(self) -> None: + inherited_signals: list[int] = [] + real_link = os.link + + def inherited_handler(signum: int, _frame: object) -> None: + inherited_signals.append(signum) + + def link_then_sigint(source: Path, target: Path, **kwargs: object) -> None: + real_link(source, target, **kwargs) + if target == self.output_path: + os.kill(os.getpid(), signal.SIGINT) + + previous_handler = signal.signal(signal.SIGINT, inherited_handler) + try: + with tempfile.TemporaryDirectory(dir=self.root) as directory: + staging = Path(directory) + staged_pte, staged_ptds = self._write_staged_artifacts(staging) + with mock.patch.object(os, "link", side_effect=link_then_sigint): + self.assertEqual( + self.receipt_path, + self._finalize(staging, staged_pte, staged_ptds), + ) + self.assertIs(signal.getsignal(signal.SIGINT), inherited_handler) + finally: + signal.signal(signal.SIGINT, previous_handler) + + self.assertEqual(inherited_signals, [signal.SIGINT]) + for path in self._published_paths(): + with self.subTest(published=path.name): + self.assertTrue(path.is_file()) + + def test_uncertain_post_link_exception_leaves_no_receipt_and_retry_fails_closed( + self, + ) -> None: + real_link = os.link + exception_injected = False + + def link_then_raise(source: Path, target: Path, **kwargs: object) -> None: + nonlocal exception_injected + real_link(source, target, **kwargs) + if target == self.output_path and not exception_injected: + exception_injected = True + raise RuntimeError("injected uncertain post-link exception") + + with tempfile.TemporaryDirectory(dir=self.root) as directory: + staging = Path(directory) + staged_pte, staged_ptds = self._write_staged_artifacts(staging) + with mock.patch.object(os, "link", side_effect=link_then_raise): + with self.assertRaisesRegex( + RuntimeError, "injected uncertain post-link exception" + ): + self._finalize(staging, staged_pte, staged_ptds) + + self.assertTrue(exception_injected) + self.assertEqual(b"mtp-pte", self.output_path.read_bytes()) + self.assertFalse(self.receipt_path.exists()) + for candidate in self._published_paths(): + if candidate != self.output_path: + self.assertFalse(candidate.exists() or candidate.is_symlink()) + + with tempfile.TemporaryDirectory(dir=self.root) as directory: + staging = Path(directory) + staged_pte, staged_ptds = self._write_staged_artifacts(staging) + with self.assertRaisesRegex(ValueError, "refusing to overwrite"): + self._finalize(staging, staged_pte, staged_ptds) + + self.assertEqual(b"mtp-pte", self.output_path.read_bytes()) + self.assertFalse(self.receipt_path.exists()) + for candidate in self._published_paths(): + if candidate != self.output_path: + self.assertFalse(candidate.exists() or candidate.is_symlink()) + + def test_preexisting_destination_is_preserved_for_each_publication(self) -> None: + for destination in self._published_paths(): + with self.subTest(destination=destination.name): + destination.parent.mkdir(parents=True, exist_ok=True) + destination.write_bytes(b"foreign-preexisting") + with tempfile.TemporaryDirectory(dir=self.root) as directory: + staging = Path(directory) + staged_pte, staged_ptds = self._write_staged_artifacts(staging) + with self.assertRaisesRegex(ValueError, "refusing to overwrite"): + self._finalize(staging, staged_pte, staged_ptds) + + self.assertEqual(b"foreign-preexisting", destination.read_bytes()) + for candidate in self._published_paths(): + if candidate != destination: + self.assertFalse(candidate.exists() or candidate.is_symlink()) + destination.unlink() + + def test_preexisting_same_inode_destination_is_not_owned_on_link_failure( + self, + ) -> None: + destination = self.artifact_root / "constants_0.ptd" + destination.parent.mkdir(parents=True) + with tempfile.TemporaryDirectory(dir=self.root) as directory: + staging = Path(directory) + staged_pte, staged_ptds = self._write_staged_artifacts(staging) + os.link(staged_ptds[0], destination) + expected_identity = destination.stat().st_dev, destination.stat().st_ino + + with self.assertRaisesRegex(ValueError, "refusing to overwrite"): + self._finalize(staging, staged_pte, staged_ptds) + + observed = destination.stat() + self.assertEqual(expected_identity, (observed.st_dev, observed.st_ino)) + self.assertEqual(b"mtp-ptd-0", destination.read_bytes()) + for candidate in self._published_paths(): + if candidate != destination: + self.assertFalse(candidate.exists() or candidate.is_symlink()) + + def test_prelink_interrupt_never_owns_same_inode_destination(self) -> None: + destination = self.artifact_root / "constants_0.ptd" + destination.parent.mkdir(parents=True) + with tempfile.TemporaryDirectory(dir=self.root) as directory: + staging = Path(directory) + staged_pte, staged_ptds = self._write_staged_artifacts(staging) + os.link(staged_ptds[0], destination) + expected_identity = destination.stat().st_dev, destination.stat().st_ino + + with mock.patch.object( + os, + "link", + side_effect=KeyboardInterrupt("injected pre-link interrupt"), + ): + with self.assertRaisesRegex( + KeyboardInterrupt, "injected pre-link interrupt" + ): + self._finalize(staging, staged_pte, staged_ptds) + + observed = destination.stat() + self.assertEqual(expected_identity, (observed.st_dev, observed.st_ino)) + self.assertEqual(b"mtp-ptd-0", destination.read_bytes()) + for candidate in self._published_paths(): + if candidate != destination: + self.assertFalse(candidate.exists() or candidate.is_symlink()) + + def test_destination_replacement_during_rollback_is_preserved(self) -> None: + destination = self.artifact_root / "constants_0.ptd" + foreign_contents = b"foreign-rollback-racer" + real_rename = os.rename + real_unlink = Path.unlink + replacement_injected = False + real_validate = gemma4_manifest.validate_mtp_manifest + validation_count = 0 + + def inject_replacement(path: Path) -> None: + nonlocal replacement_injected + if path != destination or replacement_injected: + return + replacement_injected = True + real_unlink(path) + path.write_bytes(foreign_contents) + + def rename_after_replacement( + source: Path, target: Path, *args: object, **kwargs: object + ) -> None: + inject_replacement(Path(source)) + real_rename(source, target, *args, **kwargs) + + def unlink_after_replacement( + path: Path, *args: object, **kwargs: object + ) -> None: + inject_replacement(Path(path)) + real_unlink(path, *args, **kwargs) + + def validate_then_fail(root: Path, manifest: dict[str, object]) -> None: + nonlocal validation_count + validation_count += 1 + real_validate(root, manifest) + if validation_count == 2: + raise ValueError("injected validation failure") + + with tempfile.TemporaryDirectory(dir=self.root) as directory: + staging = Path(directory) + staged_pte, staged_ptds = self._write_staged_artifacts(staging) + with ( + mock.patch.object(os, "rename", side_effect=rename_after_replacement), + mock.patch.object( + Path, "unlink", autospec=True, side_effect=unlink_after_replacement + ), + mock.patch.object( + gemma4_manifest, + "validate_mtp_manifest", + side_effect=validate_then_fail, + ), + ): + with self.assertRaisesRegex(ValueError, "injected validation failure"): + self._finalize(staging, staged_pte, staged_ptds) + + self.assertTrue(replacement_injected) + self.assertEqual(foreign_contents, destination.read_bytes()) + for candidate in self._published_paths(): + if candidate != destination: + self.assertFalse(candidate.exists() or candidate.is_symlink()) + + def test_rollback_continues_after_foreign_restore_failure(self) -> None: + earlier_destination = self.artifact_root / "constants_0.ptd" + later_destination = self.output_path + quarantined_foreign = b"foreign-moved-to-recovery" + replacement_foreign = b"foreign-later-occupant" + real_link = os.link + real_rename = os.rename + real_unlink = Path.unlink + real_validate = gemma4_manifest.validate_mtp_manifest + validation_count = 0 + replacement_injected = False + restore_blocked = False + injected_failure = ValueError("injected final validation failure") + + def rename_after_replacement( + source: Path, target: Path, *args: object, **kwargs: object + ) -> None: + nonlocal replacement_injected + if Path(source) == later_destination and not replacement_injected: + replacement_injected = True + real_unlink(later_destination) + later_destination.write_bytes(quarantined_foreign) + real_rename(source, target, *args, **kwargs) + + def block_foreign_restore(source: Path, target: Path, **kwargs: object) -> None: + nonlocal restore_blocked + if Path(target) == later_destination and Path( + source + ).parent.name.startswith(".mtp-publication-quarantine."): + restore_blocked = True + later_destination.write_bytes(replacement_foreign) + real_link(source, target, **kwargs) + + def validate_then_fail(root: Path, manifest: dict[str, object]) -> None: + nonlocal validation_count + validation_count += 1 + real_validate(root, manifest) + if validation_count == 2: + raise injected_failure + + with tempfile.TemporaryDirectory(dir=self.root) as directory: + staging = Path(directory) + staged_pte, staged_ptds = self._write_staged_artifacts(staging) + with ( + mock.patch.object(os, "rename", side_effect=rename_after_replacement), + mock.patch.object(os, "link", side_effect=block_foreign_restore), + mock.patch.object( + gemma4_manifest, + "validate_mtp_manifest", + side_effect=validate_then_fail, + ), + ): + with self.assertRaisesRegex( + ValueError, "injected final validation failure" + ) as raised: + self._finalize(staging, staged_pte, staged_ptds) + + self.assertIs(injected_failure, raised.exception) + self.assertTrue(replacement_injected) + self.assertTrue(restore_blocked) + self.assertFalse(earlier_destination.exists()) + self.assertEqual(replacement_foreign, later_destination.read_bytes()) + for candidate in self._published_paths(): + if candidate != later_destination: + self.assertFalse(candidate.exists() or candidate.is_symlink()) + + recovery_directories = [ + entry + for entry in self.artifact_root.iterdir() + if entry.is_dir() and entry.name.startswith(".mtp-publication-quarantine.") + ] + self.assertEqual(1, len(recovery_directories)) + recovery_entries = list(recovery_directories[0].iterdir()) + self.assertEqual(1, len(recovery_entries)) + self.assertEqual(quarantined_foreign, recovery_entries[0].read_bytes()) + notes = getattr(raised.exception, "__notes__", []) + self.assertEqual(1, len(notes)) + self.assertIn(f"rollback cleanup failed for {later_destination}", notes[0]) + self.assertIn("foreign publication entry retained for recovery", notes[0]) + + def test_staged_replacement_during_cleanup_is_preserved(self) -> None: + foreign_contents = b"foreign-staged-racer" + real_rename = os.rename + real_unlink = Path.unlink + replacement_injected = False + + with tempfile.TemporaryDirectory(dir=self.root) as directory: + staging = Path(directory) + staged_pte, staged_ptds = self._write_staged_artifacts(staging) + staged = staged_ptds[0] + + def inject_replacement(path: Path) -> None: + nonlocal replacement_injected + if path != staged or replacement_injected: + return + replacement_injected = True + real_unlink(path) + path.write_bytes(foreign_contents) + + def rename_after_replacement( + source: Path, target: Path, *args: object, **kwargs: object + ) -> None: + inject_replacement(Path(source)) + real_rename(source, target, *args, **kwargs) + + def unlink_after_replacement( + path: Path, *args: object, **kwargs: object + ) -> None: + inject_replacement(Path(path)) + real_unlink(path, *args, **kwargs) + + with ( + mock.patch.object(os, "rename", side_effect=rename_after_replacement), + mock.patch.object( + Path, "unlink", autospec=True, side_effect=unlink_after_replacement + ), + ): + with self.assertRaisesRegex( + ValueError, "staged artifact ownership changed" + ): + self._finalize(staging, staged_pte, staged_ptds) + + self.assertTrue(replacement_injected) + self.assertEqual(foreign_contents, staged.read_bytes()) + + self._assert_no_publications() + + def test_foreign_racer_is_preserved_before_and_after_link(self) -> None: + real_link = os.link + for timing in ("before", "after"): + for destination in self._published_paths(): + with self.subTest(timing=timing, destination=destination.name): + + def race_then_interrupt( + source: Path, + target: Path, + expected_destination: Path = destination, + race_timing: str = timing, + **kwargs: object, + ) -> None: + if target != expected_destination: + real_link(source, target, **kwargs) + return + if race_timing == "after": + real_link(source, target, **kwargs) + target.unlink() + target.write_bytes(b"foreign-racer") + raise KeyboardInterrupt(f"injected {race_timing}-link race") + + with tempfile.TemporaryDirectory(dir=self.root) as directory: + staging = Path(directory) + staged_pte, staged_ptds = self._write_staged_artifacts(staging) + with mock.patch.object( + os, + "link", + side_effect=race_then_interrupt, + ): + with self.assertRaisesRegex( + KeyboardInterrupt, f"injected {timing}-link race" + ): + self._finalize(staging, staged_pte, staged_ptds) + + self.assertEqual(b"foreign-racer", destination.read_bytes()) + for candidate in self._published_paths(): + if candidate != destination: + self.assertFalse( + candidate.exists() or candidate.is_symlink() + ) + destination.unlink() + + def test_foreign_racer_after_link_return_is_detected(self) -> None: + real_link = os.link + for destination in self._published_paths(): + with self.subTest(destination=destination.name): + + def replace_link_with_foreign( + source: Path, + target: Path, + expected_destination: Path = destination, + **kwargs: object, + ) -> None: + real_link(source, target, **kwargs) + if target == expected_destination: + target.unlink() + target.write_bytes(b"foreign-racer") + + with tempfile.TemporaryDirectory(dir=self.root) as directory: + staging = Path(directory) + staged_pte, staged_ptds = self._write_staged_artifacts(staging) + with mock.patch.object( + os, + "link", + side_effect=replace_link_with_foreign, + ): + with self.assertRaisesRegex(ValueError, "ownership changed"): + self._finalize(staging, staged_pte, staged_ptds) + + self.assertEqual(b"foreign-racer", destination.read_bytes()) + for candidate in self._published_paths(): + if candidate != destination: + self.assertFalse(candidate.exists() or candidate.is_symlink()) + destination.unlink() + + def test_publication_order_places_receipt_last(self) -> None: + observed: list[Path] = [] + real_link = os.link + + def record_link(source: Path, target: Path, **kwargs: object) -> None: + observed.append(target) + real_link(source, target, **kwargs) + + with tempfile.TemporaryDirectory(dir=self.root) as directory: + staging = Path(directory) + staged_pte, staged_ptds = self._write_staged_artifacts(staging) + with mock.patch.object(os, "link", side_effect=record_link): + self._finalize(staging, staged_pte, staged_ptds) + + self.assertEqual(list(self._published_paths()), observed) + + def test_cross_device_publication_fails_before_linking(self) -> None: + real_stat = os.stat + + def report_different_destination_device( + path: Path, *args: object, **kwargs: object + ) -> os.stat_result: + result = real_stat(path, *args, **kwargs) + if Path(path) != self.artifact_root: + return result + fields = list(result) + fields[2] = result.st_dev + 1 + return os.stat_result(fields) + + with tempfile.TemporaryDirectory(dir=self.root) as directory: + staging = Path(directory) + staged_pte, staged_ptds = self._write_staged_artifacts(staging) + with mock.patch.object( + os, + "stat", + side_effect=report_different_destination_device, + ): + with self.assertRaisesRegex(ValueError, "same filesystem"): + self._finalize(staging, staged_pte, staged_ptds) + + self._assert_no_publications() + + def test_source_receipt_basename_must_not_alias_pte_or_ptd(self) -> None: + for basename in ( + "model.pte", + "constants_0.ptd", + "constants_1.ptd", + "constants_2.ptd", + ): + with self.subTest(basename=basename): + source = self.source_root / basename + _write_source_receipt(source) + with tempfile.TemporaryDirectory(dir=self.root) as directory: + staging = Path(directory) + staged_pte, staged_ptds = self._write_staged_artifacts(staging) + with self.assertRaisesRegex( + ValueError, "duplicate normalized artifact path" + ): + self._finalize( + staging, + staged_pte, + staged_ptds, + source_receipt=source, + ) + self._assert_no_publications() + + def test_source_receipt_must_be_a_regular_non_symlink_file(self) -> None: + for name, target in ( + ("source-link.json", self.source_receipt.name), + ("dangling-source-link.json", "missing-source.json"), + ): + with self.subTest(target=target): + symlink = self.source_root / name + symlink.symlink_to(target) + with tempfile.TemporaryDirectory(dir=self.root) as directory: + staging = Path(directory) + staged_pte, staged_ptds = self._write_staged_artifacts(staging) + with self.assertRaisesRegex(ValueError, "regular non-symlink"): + self._finalize( + staging, + staged_pte, + staged_ptds, + source_receipt=symlink, + ) + self._assert_no_publications() + + def test_dangling_staged_pte_and_each_ptd_are_rejected(self) -> None: + for artifact_index in range(4): + with self.subTest(artifact_index=artifact_index): + with tempfile.TemporaryDirectory(dir=self.root) as directory: + staging = Path(directory) + staged_pte, staged_ptds = self._write_staged_artifacts(staging) + artifact = [staged_pte, *staged_ptds][artifact_index] + artifact.unlink() + artifact.symlink_to("missing-staged-artifact") + with self.assertRaisesRegex( + ValueError, "staged artifact is not regular" + ): + self._finalize(staging, staged_pte, staged_ptds) + self._assert_no_publications() + + +class CombinedRuntimeEnvelopeTest(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.plain_root = self.root / "plain-input" + self.mtp_root = self.root / "mtp-input" + self.plain_root.mkdir() + self.mtp_root.mkdir() + + plain_pte = self.plain_root / "plain.pte" + plain_source = self.plain_root / "source.json" + plain_ptds = [self.plain_root / f"plain-{index}.ptd" for index in range(3)] + plain_pte.write_bytes(b"plain-pte") + _write_source_receipt(plain_source) + for index, path in enumerate(plain_ptds): + path.write_bytes(f"plain-{index}".encode()) + self.plain_manifest = gemma4_manifest.create_plain_manifest( + self.plain_root, + {"pte": Path(plain_pte.name), "source": Path(plain_source.name)}, + [Path(path.name) for path in plain_ptds], + ) + + self.mtp_pte = self.mtp_root / "mtp.pte" + self.mtp_source = self.mtp_root / "source.json" + self.mtp_ptds = [self.mtp_root / f"mtp-{index}.ptd" for index in range(3)] + self.mtp_pte.write_bytes(b"mtp-pte") + _write_source_receipt(self.mtp_source) + for index, path in enumerate(self.mtp_ptds): + path.write_bytes(f"mtp-{index}".encode()) + self.mtp_manifest = gemma4_manifest.create_mtp_manifest( + self.mtp_root, + { + "pte": Path(self.mtp_pte.name), + "source": Path(self.mtp_source.name), + }, + [Path(path.name) for path in self.mtp_ptds], + ) + self.mtp_manifest.update( + { + "evidence": _generated_mtp_evidence(), + "provenance": copy.deepcopy( + gemma4_manifest.MTP_SOURCE_VERIFIED_PROVENANCE + ), + } + ) + + self.plain_receipt = self.root / "plain.json" + self.mtp_receipt = self.root / "mtp.json" + self.runtime_receipt = self.root / "runtime-source.json" + self.plain_receipt.write_text(json.dumps(self.plain_manifest)) + self.mtp_receipt.write_text(json.dumps(self.mtp_manifest)) + self.runtime_paths: dict[str, dict[str, Path]] = { + "profile": { + "javascript": self.root / "mtp-profile/gemma4_mtp_profile.js", + "wasm": self.root / "mtp-profile/gemma4_mtp_profile.wasm", + }, + "wall": { + "javascript": self.root / "mtp-wall/gemma4_mtp.js", + "wasm": self.root / "mtp-wall/gemma4_mtp.wasm", + }, + } + for flavor, paths in self.runtime_paths.items(): + for kind, path in paths.items(): + path.parent.mkdir(exist_ok=True) + path.write_bytes(f"{flavor}-{kind}".encode()) + self.plain_runtime_paths: dict[str, dict[str, Path]] = { + "profile": { + "javascript": self.root / "plain-profile/webgpu_llama.js", + "wasm": self.root / "plain-profile/webgpu_llama.wasm", + }, + "wall": { + "javascript": self.root / "plain-wall/webgpu_llama.js", + "wasm": self.root / "plain-wall/webgpu_llama.wasm", + }, + } + for flavor, paths in self.plain_runtime_paths.items(): + for kind, path in paths.items(): + path.parent.mkdir(exist_ok=True) + value = "plain-javascript" if kind == "javascript" else flavor + path.write_bytes(value.encode()) + self.source_manifest = self.root / "source-manifest.json" + self.wgsl_manifest = self.root / "wgsl-manifest.json" + self.source_manifest.write_text(json.dumps(_test_source_manifest())) + self.wgsl_manifest.write_text(json.dumps(_test_wgsl_manifest())) + self.build_commands = { + "mtp": { + "profile": self.root / "mtp-profile-recipe.json", + "wall": self.root / "mtp-wall-recipe.json", + }, + "plain": { + "profile": self.root / "plain-profile-recipe.json", + "wall": self.root / "plain-wall-recipe.json", + }, + } + for model, builds in self.build_commands.items(): + for flavor, path in builds.items(): + path.write_text( + json.dumps( + gemma4_manifest.canonical_build_recipe(model, flavor), + indent=2, + sort_keys=True, + ) + + "\n" + ) + self._write_runtime_receipt() + self.target_prefill_receipt = self.root / "target-prefill.json" + self._write_target_prefill_receipt() + self.destination = self.root / "staged" + + def _write_runtime_receipt(self) -> None: + receipt = self._create_runtime_receipt() + self.runtime_receipt.write_text(json.dumps(receipt)) + + def _write_target_prefill_receipt(self) -> None: + self.target_prefill_receipt.write_text( + json.dumps(_target_prefill_receipt(self.runtime_receipt)), + encoding="utf-8", + ) + + def _create_runtime_receipt( + self, + *, + runtime_paths: dict[str, object] | None = None, + build_command_paths: dict[str, object] | None = None, + ) -> dict[str, object]: + with mock.patch.object( + gemma4_manifest, + "create_source_manifest", + return_value=_test_source_manifest(), + ), mock.patch.object( + gemma4_manifest, + "create_wgsl_manifest", + return_value=_test_wgsl_manifest(), + ): + return gemma4_manifest.create_runtime_source_receipt( + fbsource_root=self.root, + oss_root=self.root, + backend_root=self.root, + source_manifest_path=self.source_manifest, + wgsl_manifest_path=self.wgsl_manifest, + manifest_paths={ + "mtp": self.mtp_receipt, + "plain": self.plain_receipt, + }, + model_roots={"mtp": self.mtp_root, "plain": self.plain_root}, + runtime_paths=runtime_paths + or { + "mtp": self.runtime_paths, + "plain": self.plain_runtime_paths, + }, + build_command_paths=build_command_paths or self.build_commands, + ) + + def _demote_mtp_to_pending(self) -> dict[str, object]: + """Pending means no source receipt at all, not merely a pending label.""" + self.mtp_source.unlink() + manifest = gemma4_manifest.create_mtp_manifest( + self.mtp_root, + {"pte": Path(self.mtp_pte.name)}, + [Path(path.name) for path in self.mtp_ptds], + ) + manifest["evidence"] = _generated_mtp_evidence() + return manifest + + def _staged_envelope(self) -> dict[str, object]: + return json.loads( + (self.destination / "gemma4_webgpu_combined_runtime.json").read_text() + ) + + def _stage(self, *, refresh_target_prefill: bool = True) -> None: + if refresh_target_prefill: + self._write_target_prefill_receipt() + gemma4_manifest.stage_combined_runtime( + self.destination, + self.plain_root, + self.plain_receipt, + self.mtp_root, + self.mtp_receipt, + self.runtime_receipt, + self.runtime_paths["wall"]["javascript"], + self.runtime_paths["wall"]["wasm"], + self.runtime_paths["profile"]["javascript"], + self.runtime_paths["profile"]["wasm"], + plain_profile_javascript_path=self.plain_runtime_paths["profile"][ + "javascript" + ], + plain_profile_wasm_path=self.plain_runtime_paths["profile"]["wasm"], + plain_wall_javascript_path=self.plain_runtime_paths["wall"]["javascript"], + plain_wall_wasm_path=self.plain_runtime_paths["wall"]["wasm"], + source_manifest_path=self.source_manifest, + wgsl_manifest_path=self.wgsl_manifest, + build_recipe_paths=self.build_commands, + target_prefill_receipt_path=self.target_prefill_receipt, + ) + + def test_target_prefill_receipt_is_staged_in_envelope_v3(self) -> None: + self._stage() + envelope = self._staged_envelope() + self.assertEqual(envelope["schema_version"], 3) + self.assertEqual(set(envelope["receipts"]), {"mtp", "plain", "target_prefill"}) + identity = envelope["receipts"]["target_prefill"] + self.assertEqual(identity["path"], "receipts/target_prefill.json") + self.assertEqual( + {key: identity[key] for key in ("bytes", "sha256")}, + _identity(self.target_prefill_receipt), + ) + + def test_final_schema_versions_are_split_by_receipt_role(self) -> None: + self.assertEqual(_sealed_source_receipt()["schema_version"], 3) + self.assertEqual( + gemma4_manifest.canonical_build_recipe("plain", "profile")[ + "schema_version" + ], + 2, + ) + self.assertEqual( + json.loads(self.runtime_receipt.read_text())["schema_version"], 4 + ) + self._stage() + self.assertEqual(self._staged_envelope()["schema_version"], 3) + + def test_build_recipes_bind_model_specific_factories_and_stems(self) -> None: + expected = { + ("mtp", "profile"): ("createGemma4MtpProfile", "gemma4_mtp_profile"), + ("mtp", "wall"): ("createGemma4Mtp", "gemma4_mtp"), + ("plain", "profile"): ("createWebGPULlama", "webgpu_llama"), + ("plain", "wall"): ("createWebGPULlama", "webgpu_llama"), + } + for (model, flavor), (factory, output_stem) in expected.items(): + with self.subTest(model=model, flavor=flavor): + recipe = gemma4_manifest.canonical_build_recipe(model, flavor) + self.assertEqual(recipe["factory"], factory) + self.assertEqual(recipe["output_stem"], output_stem) + self.assertEqual(recipe["profiling_enabled"], flavor == "profile") + self.assertTrue( + str(recipe["outputs"]["javascript"]).endswith(f"/{output_stem}.js") + ) + self.assertTrue( + str(recipe["outputs"]["wasm"]).endswith(f"/{output_stem}.wasm") + ) + + def test_combined_staging_requires_target_prefill_receipt(self) -> None: + self.assertIn( + "target_prefill_receipt_path", + inspect.signature(gemma4_manifest.stage_combined_runtime).parameters, + ) + + def test_target_prefill_receipt_must_bind_runtime_source(self) -> None: + receipt = json.loads(self.target_prefill_receipt.read_text()) + receipt["producer"]["runtime_source_receipt"]["sha256"] = "f" * 64 + self.target_prefill_receipt.write_text(json.dumps(receipt)) + with self.assertRaisesRegex(ValueError, "runtime source identity"): + self._stage(refresh_target_prefill=False) + + def test_target_prefill_receipt_must_bind_runtime_source_head(self) -> None: + receipt = json.loads(self.target_prefill_receipt.read_text()) + receipt["producer"]["fbsource_commit"] = "f" * 40 + self.target_prefill_receipt.write_text(json.dumps(receipt)) + with self.assertRaisesRegex(ValueError, "fbsource commit"): + self._stage(refresh_target_prefill=False) + + def test_target_prefill_receipt_tamper_after_staging_is_rejected(self) -> None: + self._stage() + envelope = self._staged_envelope() + path = self.destination / "receipts/target_prefill.json" + receipt = json.loads(path.read_text()) + receipt["contexts"]["513"]["prefill_token_raw"] += 1 + path.write_text(json.dumps(receipt)) + with self.assertRaisesRegex(ValueError, "target.prefill|target-prefill"): + gemma4_manifest.validate_combined_runtime_envelope( + self.destination, envelope + ) + + def test_recanonicalized_target_receipt_swap_still_rechecks_source_link( + self, + ) -> None: + self._stage() + envelope = self._staged_envelope() + target_path = self.destination / "receipts/target_prefill.json" + swapped = json.loads(target_path.read_text()) + other_runtime_source = self.root / "other-runtime-source.json" + other_runtime_source.write_bytes( + (self.destination / "receipts/runtime_source.json").read_bytes() + b"\n" + ) + swapped["producer"]["runtime_source_receipt"] = _identity(other_runtime_source) + target_path.write_text(json.dumps(swapped), encoding="utf-8") + envelope["receipts"]["target_prefill"].update(_identity(target_path)) + + with self.assertRaisesRegex(ValueError, "runtime source identity"): + gemma4_manifest.validate_combined_runtime_envelope( + self.destination, envelope + ) + + def test_recanonicalized_target_receipt_still_binds_producer_resource( + self, + ) -> None: + self._stage() + envelope = self._staged_envelope() + target_path = self.destination / "receipts/target_prefill.json" + swapped = json.loads(target_path.read_text()) + swapped["producer"]["source_sha256"] = "f" * 64 + target_path.write_text(json.dumps(swapped), encoding="utf-8") + envelope["receipts"]["target_prefill"].update(_identity(target_path)) + + with self.assertRaisesRegex(ValueError, "producer source hash"): + gemma4_manifest.validate_combined_runtime_envelope( + self.destination, envelope + ) + + def test_combined_runtime_v2_is_rejected(self) -> None: + self._stage() + envelope = self._staged_envelope() + envelope["schema_version"] = 2 + with self.assertRaisesRegex(ValueError, "schema version"): + gemma4_manifest.validate_combined_runtime_envelope( + self.destination, envelope + ) + + def test_plain_wall_and_profile_runtimes_are_staged_and_bound(self) -> None: + self._stage() + plain = self._staged_envelope()["runtime"]["plain"] + self.assertEqual(set(plain), {"profile", "wall"}) + for flavor, paths in self.plain_runtime_paths.items(): + for kind, suffix in (("javascript", "js"), ("wasm", "wasm")): + with self.subTest(flavor=flavor, kind=kind): + identity = plain[flavor][kind] + self.assertEqual(sorted(identity), ["bytes", "path", "sha256"]) + self.assertEqual( + identity["path"], f"runtime/plain/{flavor}.{suffix}" + ) + self.assertEqual( + {key: identity[key] for key in ("bytes", "sha256")}, + _identity(paths[kind]), + ) + self.assertEqual( + (self.destination / str(identity["path"])).read_bytes(), + paths[kind].read_bytes(), + ) + + def test_plain_identical_javascript_still_binds_distinct_wasm_pairs( + self, + ) -> None: + receipt = self._create_runtime_receipt() + wall = receipt["runtime"]["plain"]["wall"] + profile = receipt["runtime"]["plain"]["profile"] + self.assertEqual(wall["javascript"], profile["javascript"]) + self.assertNotEqual(wall["wasm"], profile["wasm"]) + + def test_plain_runtime_removes_the_missing_adapter_status(self) -> None: + self._stage() + views = self._staged_envelope()["views"] + self.assertNotIn("blocked_missing_generic_browser_adapter", json.dumps(views)) + self.assertEqual(views["plain"]["runtime"], "plain") + + def test_tampered_plain_runtime_is_rejected(self) -> None: + self._stage() + envelope = self._staged_envelope() + wasm = envelope["runtime"]["plain"]["wall"]["wasm"] + (self.destination / str(wasm["path"])).write_bytes(b"tampered") + with self.assertRaisesRegex(ValueError, "plain wall wasm is not bound"): + gemma4_manifest.validate_combined_runtime_envelope( + self.destination, envelope + ) + + def test_runtime_receipt_binds_plain_wall_bytes(self) -> None: + self.plain_runtime_paths["wall"]["wasm"].write_bytes(b"rebuilt-after-receipt") + with self.assertRaisesRegex(ValueError, "plain wall wasm is not bound"): + self._stage() + + def test_runtime_receipt_v1_is_rejected(self) -> None: + receipt = json.loads(self.runtime_receipt.read_text()) + receipt["schema_version"] = 1 + self.runtime_receipt.write_text(json.dumps(receipt)) + with self.assertRaisesRegex(ValueError, "schema version 4"): + self._stage() + + def test_runtime_receipt_v2_is_rejected(self) -> None: + receipt = json.loads(self.runtime_receipt.read_text()) + receipt["schema_version"] = 2 + self.runtime_receipt.write_text(json.dumps(receipt)) + with self.assertRaisesRegex(ValueError, "schema version 4"): + self._stage() + + def test_runtime_receipt_v3_is_rejected(self) -> None: + receipt = json.loads(self.runtime_receipt.read_text()) + receipt["schema_version"] = 3 + self.runtime_receipt.write_text(json.dumps(receipt)) + with self.assertRaisesRegex(ValueError, "schema version 4"): + self._stage() + + def test_runtime_receipt_does_not_accept_caller_supplied_heads(self) -> None: + parameters = inspect.signature( + gemma4_manifest.create_runtime_source_receipt + ).parameters + self.assertNotIn("fbsource_commit", parameters) + self.assertNotIn("oss_commit", parameters) + + def test_runtime_receipt_rejects_model_source_receipt_from_other_head( + self, + ) -> None: + receipt = _sealed_source_receipt() + receipt["fbsource_commit"] = "9" * 40 + receipt["oss_commit"] = "8" * 40 + source_manifest = receipt["source_manifest"] + wgsl_manifest = receipt["wgsl_manifest"] + assert isinstance(source_manifest, dict) and isinstance(wgsl_manifest, dict) + checkouts = source_manifest["checkouts"] + assert isinstance(checkouts, dict) + checkouts["fbsource"]["head"] = "9" * 40 + checkouts["oss"]["head"] = "8" * 40 + wgsl_manifest["fbsource_commit"] = "9" * 40 + (self.plain_root / "source.json").write_text(json.dumps(receipt)) + self.mtp_source.write_text(json.dumps(receipt)) + + self.plain_manifest = gemma4_manifest.create_plain_manifest( + self.plain_root, + {"pte": Path("plain.pte"), "source": Path("source.json")}, + [Path(f"plain-{index}.ptd") for index in range(3)], + ) + self.mtp_manifest = gemma4_manifest.create_mtp_manifest( + self.mtp_root, + {"pte": Path("mtp.pte"), "source": Path("source.json")}, + [Path(f"mtp-{index}.ptd") for index in range(3)], + ) + self.mtp_manifest["evidence"] = _generated_mtp_evidence() + self.plain_receipt.write_text(json.dumps(self.plain_manifest)) + self.mtp_receipt.write_text(json.dumps(self.mtp_manifest)) + + with self.assertRaisesRegex(ValueError, "source receipt.*head"): + self._write_runtime_receipt() + + def test_runtime_receipt_rejects_opaque_source_manifest(self) -> None: + self.source_manifest.write_text('{"source": "current"}') + with self.assertRaisesRegex(ValueError, "source manifest"): + self._write_runtime_receipt() + + def test_runtime_receipt_rejects_opaque_wgsl_manifest(self) -> None: + self.wgsl_manifest.write_text('{"wgsl": "current"}') + with self.assertRaisesRegex(ValueError, "WGSL manifest"): + self._write_runtime_receipt() + + def test_runtime_receipt_rejects_text_build_recipe(self) -> None: + self.build_commands["plain"]["wall"].write_text("build plain wall\n") + with self.assertRaisesRegex(ValueError, "plain wall recipe"): + self._write_runtime_receipt() + + def test_build_recipe_v1_is_rejected(self) -> None: + path = self.build_commands["plain"]["wall"] + recipe = json.loads(path.read_text()) + recipe["schema_version"] = 1 + path.write_text(json.dumps(recipe)) + with self.assertRaisesRegex(ValueError, "plain wall recipe"): + self._write_runtime_receipt() + + def test_runtime_receipt_rejects_source_copy_identity_mutation(self) -> None: + source = json.loads(self.source_manifest.read_text()) + for copy_identity in source["files"][0]["copies"].values(): + copy_identity["sha256"] = "0" * 64 + self.source_manifest.write_text(json.dumps(source)) + with self.assertRaisesRegex(ValueError, "live clean checkouts"): + self._write_runtime_receipt() + + def test_runtime_receipt_rejects_structurally_valid_wgsl_hash_mutation( + self, + ) -> None: + manifest = json.loads(self.wgsl_manifest.read_text()) + manifest["files"][0]["sha256"] = "0" * 64 + self.wgsl_manifest.write_text(json.dumps(manifest)) + with self.assertRaisesRegex(ValueError, "live generator closure"): + self._write_runtime_receipt() + + def test_runtime_receipt_rejects_missing_wgsl_registry(self) -> None: + manifest = json.loads(self.wgsl_manifest.read_text()) + manifest["files"] = [ + entry for entry in manifest["files"] if entry["role"] != "global_registry" + ] + manifest["file_set_sha256"] = _set_digest( + [ + {"path": entry["path"], "role": entry["role"]} + for entry in manifest["files"] + ] + ) + self.wgsl_manifest.write_text(json.dumps(manifest)) + with self.assertRaisesRegex(ValueError, "WGSL manifest.*registry"): + self._write_runtime_receipt() + + def test_runtime_receipt_rejects_declared_wgsl_orphan(self) -> None: + manifest = json.loads(self.wgsl_manifest.read_text()) + manifest["orphans"] = ["runtime/ops/add/orphan_wgsl.h"] + self.wgsl_manifest.write_text(json.dumps(manifest)) + with self.assertRaisesRegex(ValueError, "WGSL manifest.*orphan"): + self._write_runtime_receipt() + + def test_runtime_receipt_rejects_recipe_contract_mutations(self) -> None: + mutations = { + "target": ("target", "wrong_target"), + "profile": ("profiling_enabled", True), + "factory": ("factory", "createWrongFactory"), + } + for label, (key, value) in mutations.items(): + with self.subTest(label=label): + path = self.build_commands["plain"]["wall"] + original = path.read_text() + recipe = json.loads(original) + recipe[key] = value + path.write_text(json.dumps(recipe)) + with self.assertRaisesRegex(ValueError, "plain wall recipe"): + self._write_runtime_receipt() + path.write_text(original) + + path = self.build_commands["mtp"]["wall"] + recipe = json.loads(path.read_text()) + recipe["outputs"]["wasm"] = "wrong.wasm" + path.write_text(json.dumps(recipe)) + with self.assertRaisesRegex(ValueError, "MTP wall recipe"): + self._write_runtime_receipt() + + def test_runtime_receipt_rejects_factory_or_output_stem_mutation(self) -> None: + for key, value in ( + ("factory", "createWrongFactory"), + ("output_stem", "wrong_output"), + ): + with self.subTest(key=key): + receipt = json.loads(self.runtime_receipt.read_text()) + receipt["runtime"]["mtp"]["profile"][key] = value + self.runtime_receipt.write_text(json.dumps(receipt)) + with self.assertRaisesRegex(ValueError, f"{key} mismatch"): + self._stage() + self._write_runtime_receipt() + + def test_runtime_receipt_rejects_noncanonical_product_basename(self) -> None: + runtime_paths = { + "mtp": { + flavor: dict(paths) for flavor, paths in self.runtime_paths.items() + }, + "plain": { + flavor: dict(paths) + for flavor, paths in self.plain_runtime_paths.items() + }, + } + wrong = self.root / "mtp-wall/wrong.js" + wrong.write_bytes(b"wrong-basename") + runtime_paths["mtp"]["wall"]["javascript"] = wrong + with self.assertRaisesRegex(ValueError, "JavaScript basename mismatch"): + self._create_runtime_receipt(runtime_paths=runtime_paths) + + def test_runtime_receipt_records_only_validated_not_executed_claims(self) -> None: + receipt = json.loads(self.runtime_receipt.read_text()) + self.assertEqual( + receipt.get("verification"), + { + "build_execution": "not_attested", + "recipe": "validated", + "source_checkout": "verified", + "wgsl_codegen": "verified", + }, + ) + + def test_runtime_receipt_rejects_unattested_execution_claim_upgrade(self) -> None: + receipt = json.loads(self.runtime_receipt.read_text()) + receipt["verification"]["build_execution"] = "verified" + self.runtime_receipt.write_text(json.dumps(receipt)) + with self.assertRaisesRegex(ValueError, "verification claims"): + self._stage() + + def test_runtime_receipt_generator_binds_source_builds_and_artifacts( + self, + ) -> None: + receipt = self._create_runtime_receipt() + self.assertEqual(receipt["schema_version"], 4) + self.assertEqual(receipt["fbsource_commit"], "1" * 40) + self.assertEqual(receipt["oss_commit"], "2" * 40) + self.assertEqual( + receipt["source_inputs"]["source_manifest"], + { + "path": "closure/source_manifest.json", + **_identity(self.source_manifest), + }, + ) + self.assertEqual( + receipt["source_inputs"]["wgsl_manifest"], + { + "path": "closure/wgsl_manifest.json", + **_identity(self.wgsl_manifest), + }, + ) + self.assertEqual(receipt["runtime"]["plain"]["target"], "gemma4_plain_wasm") + self.assertEqual(receipt["runtime"]["mtp"]["target"], "gemma4_spec_browser") + self.assertEqual( + receipt["runtime"]["plain"]["profile"]["factory"], + "createWebGPULlama", + ) + self.assertEqual( + receipt["runtime"]["mtp"]["profile"]["factory"], + "createGemma4MtpProfile", + ) + self.assertEqual( + receipt["runtime"]["mtp"]["profile"]["output_stem"], + "gemma4_mtp_profile", + ) + self.assertEqual( + receipt["runtime"]["plain"]["wall"]["recipe"], + { + "path": "closure/recipes/plain-wall.json", + **_identity(self.build_commands["plain"]["wall"]), + }, + ) + + def test_runtime_receipt_rejects_aliased_mtp_roles(self) -> None: + for role in ("javascript", "wasm", "recipe"): + with self.subTest(role=role): + runtime_paths = { + "mtp": { + flavor: dict(paths) + for flavor, paths in self.runtime_paths.items() + }, + "plain": { + flavor: dict(paths) + for flavor, paths in self.plain_runtime_paths.items() + }, + } + build_recipes = { + model: dict(paths) for model, paths in self.build_commands.items() + } + if role == "recipe": + build_recipes["mtp"]["profile"] = build_recipes["mtp"]["wall"] + else: + runtime_paths["mtp"]["profile"][role] = runtime_paths["mtp"][ + "wall" + ][role] + with self.assertRaisesRegex( + ValueError, f"wall/profile {role} identities must differ" + ): + self._create_runtime_receipt( + runtime_paths=runtime_paths, + build_command_paths=build_recipes, + ) + + def test_runtime_receipt_rejects_aliased_plain_wasm_recipe_or_pair( + self, + ) -> None: + for role in ("pair", "wasm", "recipe"): + with self.subTest(role=role): + runtime_paths = { + "mtp": { + flavor: dict(paths) + for flavor, paths in self.runtime_paths.items() + }, + "plain": { + flavor: dict(paths) + for flavor, paths in self.plain_runtime_paths.items() + }, + } + build_recipes = { + model: dict(paths) for model, paths in self.build_commands.items() + } + if role == "recipe": + build_recipes["plain"]["profile"] = build_recipes["plain"]["wall"] + else: + runtime_paths["plain"]["profile"]["wasm"] = runtime_paths["plain"][ + "wall" + ]["wasm"] + if role == "pair": + runtime_paths["plain"]["profile"]["javascript"] = runtime_paths[ + "plain" + ]["wall"]["javascript"] + else: + distinct_javascript = ( + self.root / "plain-profile-distinct/webgpu_llama.js" + ) + distinct_javascript.parent.mkdir(exist_ok=True) + distinct_javascript.write_bytes(b"distinct-javascript") + runtime_paths["plain"]["profile"][ + "javascript" + ] = distinct_javascript + with self.assertRaisesRegex( + ValueError, f"wall/profile {role} identities must differ" + ): + self._create_runtime_receipt( + runtime_paths=runtime_paths, + build_command_paths=build_recipes, + ) + + def test_staging_rehashes_source_and_build_recipe_bytes(self) -> None: + for label, path in ( + ("source manifest", self.source_manifest), + ("WGSL manifest", self.wgsl_manifest), + ("plain profile recipe", self.build_commands["plain"]["profile"]), + ("plain wall recipe", self.build_commands["plain"]["wall"]), + ("MTP wall recipe", self.build_commands["mtp"]["wall"]), + ("MTP profile recipe", self.build_commands["mtp"]["profile"]), + ): + with self.subTest(label=label): + original = path.read_bytes() + path.write_bytes(original + b"tampered") + with self.assertRaisesRegex(ValueError, "is not bound"): + self._stage() + path.write_bytes(original) + + def test_staged_closure_carries_every_bound_source_and_recipe_byte(self) -> None: + self._stage() + expected = { + "closure/source_manifest.json": self.source_manifest, + "closure/wgsl_manifest.json": self.wgsl_manifest, + "closure/recipes/plain-profile.json": self.build_commands["plain"][ + "profile" + ], + "closure/recipes/plain-wall.json": self.build_commands["plain"]["wall"], + "closure/recipes/mtp-wall.json": self.build_commands["mtp"]["wall"], + "closure/recipes/mtp-profile.json": self.build_commands["mtp"]["profile"], + } + for relative, source in expected.items(): + with self.subTest(relative=relative): + self.assertEqual( + (self.destination / relative).read_bytes(), source.read_bytes() + ) + + def test_pending_plain_provenance_cannot_stage(self) -> None: + self.plain_manifest["provenance"] = copy.deepcopy( + gemma4_manifest.MTP_ACCEPTED_PROVENANCE + ) + self.plain_receipt.write_text(json.dumps(self.plain_manifest)) + with self.assertRaisesRegex(ValueError, "plain.*provenance"): + self._write_runtime_receipt() + + def test_stages_and_validates_exact_runtime_bytes(self) -> None: + self._stage() + manifest_path = self.destination / "gemma4_webgpu_combined_runtime.json" + envelope = json.loads(manifest_path.read_text()) + gemma4_manifest.validate_combined_runtime_envelope(self.destination, envelope) + self.assertEqual( + envelope["views"]["mtp"]["status"], + "pending_gpu_execution_validation", + ) + + def test_create_runtime_source_cli_carries_plain_profile_inputs(self) -> None: + output = self.root / "runtime-source-cli.json" + argv = [ + "create-runtime-source", + "--output", + str(output), + "--fbsource-root", + str(self.root), + "--oss-root", + str(self.root), + "--backend-root", + str(self.root), + "--source-manifest", + str(self.source_manifest), + "--wgsl-manifest", + str(self.wgsl_manifest), + "--plain-manifest", + str(self.plain_receipt), + "--mtp-manifest", + str(self.mtp_receipt), + "--plain-root", + str(self.plain_root), + "--mtp-root", + str(self.mtp_root), + ] + for model, paths in ( + ("plain", self.plain_runtime_paths), + ("mtp", self.runtime_paths), + ): + for flavor, artifacts in paths.items(): + argv.extend( + [ + f"--{model}-{flavor}-javascript", + str(artifacts["javascript"]), + f"--{model}-{flavor}-wasm", + str(artifacts["wasm"]), + f"--{model}-{flavor}-recipe", + str(self.build_commands[model][flavor]), + ] + ) + with mock.patch.object( + gemma4_manifest, + "create_source_manifest", + return_value=_test_source_manifest(), + ), mock.patch.object( + gemma4_manifest, + "create_wgsl_manifest", + return_value=_test_wgsl_manifest(), + ): + self.assertEqual(gemma4_manifest.main(argv), 0) + receipt = json.loads(output.read_text()) + self.assertEqual( + set(receipt["runtime"]["plain"]), {"profile", "target", "wall"} + ) + + def test_stage_runtime_cli_carries_every_plain_profile_input(self) -> None: + self._write_target_prefill_receipt() + destination = self.root / "staged-cli" + argv = [ + "stage-runtime", + "--destination-root", + str(destination), + "--plain-root", + str(self.plain_root), + "--plain-receipt", + str(self.plain_receipt), + "--mtp-root", + str(self.mtp_root), + "--mtp-receipt", + str(self.mtp_receipt), + "--runtime-source-receipt", + str(self.runtime_receipt), + "--target-prefill-receipt", + str(self.target_prefill_receipt), + "--plain-wall-javascript", + str(self.plain_runtime_paths["wall"]["javascript"]), + "--plain-wall-wasm", + str(self.plain_runtime_paths["wall"]["wasm"]), + "--plain-wall-recipe", + str(self.build_commands["plain"]["wall"]), + "--plain-profile-javascript", + str(self.plain_runtime_paths["profile"]["javascript"]), + "--plain-profile-wasm", + str(self.plain_runtime_paths["profile"]["wasm"]), + "--plain-profile-recipe", + str(self.build_commands["plain"]["profile"]), + "--mtp-wall-javascript", + str(self.runtime_paths["wall"]["javascript"]), + "--mtp-wall-wasm", + str(self.runtime_paths["wall"]["wasm"]), + "--mtp-wall-recipe", + str(self.build_commands["mtp"]["wall"]), + "--mtp-profile-javascript", + str(self.runtime_paths["profile"]["javascript"]), + "--mtp-profile-wasm", + str(self.runtime_paths["profile"]["wasm"]), + "--mtp-profile-recipe", + str(self.build_commands["mtp"]["profile"]), + "--source-manifest", + str(self.source_manifest), + "--wgsl-manifest", + str(self.wgsl_manifest), + ] + self.assertEqual(gemma4_manifest.main(argv), 0) + for relative, source in ( + ( + "runtime/plain/profile.js", + self.plain_runtime_paths["profile"]["javascript"], + ), + ( + "runtime/plain/profile.wasm", + self.plain_runtime_paths["profile"]["wasm"], + ), + ( + "closure/recipes/plain-profile.json", + self.build_commands["plain"]["profile"], + ), + ): + with self.subTest(relative=relative): + self.assertEqual( + (destination / relative).read_bytes(), source.read_bytes() + ) + + def test_rejects_accepted_oracle_as_current_source(self) -> None: + self.mtp_manifest = self._demote_mtp_to_pending() + self.mtp_manifest["provenance"] = copy.deepcopy( + gemma4_manifest.MTP_ACCEPTED_PROVENANCE + ) + del self.mtp_manifest["evidence"] + self.mtp_receipt.write_text(json.dumps(self.mtp_manifest)) + with self.assertRaisesRegex(ValueError, "source closure is still pending"): + self._write_runtime_receipt() + + def test_rejects_pending_source_mtp_manifest(self) -> None: + self.mtp_receipt.write_text(json.dumps(self._demote_mtp_to_pending())) + with self.assertRaisesRegex(ValueError, "source closure is still pending"): + self._write_runtime_receipt() + + def test_sealed_source_receipt_stages_as_source_verified(self) -> None: + self.assertEqual( + self.mtp_manifest["provenance"], + gemma4_manifest.MTP_SOURCE_VERIFIED_PROVENANCE, + ) + self._stage() + gemma4_manifest.validate_combined_runtime_envelope( + self.destination, self._staged_envelope() + ) + + def test_status_only_pending_label_cannot_stage(self) -> None: + self._stage() + envelope = self._staged_envelope() + envelope["source_verification"]["mtp"]["provenance"][ + "source_closure" + ] = "pending_D10_reproduction_receipt" + with self.assertRaisesRegex(ValueError, "source closure is still pending"): + gemma4_manifest.validate_combined_runtime_envelope( + self.destination, envelope + ) + + def test_envelope_carries_a_source_verification_block(self) -> None: + # Source closure only; PTE/WASM behaviour is a separate runtime gate. + self._stage() + verification = self._staged_envelope()["source_verification"] + self.assertEqual(sorted(verification), ["mtp", "plain"]) + for label in ("mtp", "plain"): + with self.subTest(view=label): + self.assertEqual( + sorted(verification[label]), ["provenance", "source_receipt"] + ) + self.assertEqual( + sorted(verification[label]["source_receipt"]), + ["bytes", "path", "sha256"], + ) + self.assertEqual( + verification[label]["source_receipt"]["path"], "source.json" + ) + self.assertEqual( + verification["mtp"]["provenance"], + gemma4_manifest.MTP_SOURCE_VERIFIED_PROVENANCE, + ) + self.assertIsNone(verification["plain"]["provenance"]) + + def test_source_verification_is_independent_of_runtime_evidence(self) -> None: + self._stage() + first = self._staged_envelope() + for flavor, paths in self.runtime_paths.items(): + for kind, path in paths.items(): + path.write_bytes(f"{flavor}-{kind}-rebuilt".encode()) + self._write_runtime_receipt() + self.destination = self.root / "staged-rebuilt" + self._stage() + second = self._staged_envelope() + self.assertNotEqual(first["runtime"], second["runtime"]) + self.assertEqual(first["source_verification"], second["source_verification"]) + + def test_tampered_source_receipt_is_not_masked_by_valid_runtime(self) -> None: + self._stage() + envelope = self._staged_envelope() + (self.destination / "mtp" / "source.json").write_text( + json.dumps({**_SEALED_SOURCE_RECEIPT, "oss_commit": "3" * 40}) + ) + wasm = envelope["runtime"]["mtp"]["wall"]["wasm"] + self.assertEqual( + _identity(self.destination / str(wasm["path"])), + {"bytes": wasm["bytes"], "sha256": wasm["sha256"]}, + ) + with self.assertRaisesRegex( + ValueError, "mtp source receipt byte or SHA-256 identity mismatch" + ): + gemma4_manifest.validate_combined_runtime_envelope( + self.destination, envelope + ) + + def test_rejects_runtime_not_bound_to_source_receipt(self) -> None: + self.runtime_paths["wall"]["wasm"].write_bytes(b"tampered") + with self.assertRaisesRegex(ValueError, "bound to its build receipt"): + self._stage() + + def test_rejects_extra_staged_file(self) -> None: + self._stage() + (self.destination / "extra.bin").write_bytes(b"extra") + envelope = json.loads( + (self.destination / "gemma4_webgpu_combined_runtime.json").read_text() + ) + with self.assertRaisesRegex(ValueError, "extra files"): + gemma4_manifest.validate_combined_runtime_envelope( + self.destination, envelope + ) + + def test_rejects_swapped_runtime_roles(self) -> None: + self._stage() + envelope = json.loads( + (self.destination / "gemma4_webgpu_combined_runtime.json").read_text() + ) + profile = envelope["runtime"]["mtp"]["profile"] + wall = envelope["runtime"]["mtp"]["wall"] + profile["javascript"], wall["javascript"] = ( + wall["javascript"], + profile["javascript"], + ) + with self.assertRaisesRegex(ValueError, "non-canonical role bindings"): + gemma4_manifest.validate_combined_runtime_envelope( + self.destination, envelope + ) + + +class PlainManifestContractTest(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.source_receipt = self.root / "source_receipt.json" + self.ptds = [self.root / f"constants_{index}.ptd" for index in range(3)] + self.pte.write_bytes(b"plain-pte") + _write_source_receipt(self.source_receipt) + for index, path in enumerate(self.ptds): + path.write_bytes(f"plain-ptd-{index}".encode("utf-8")) + self.manifest = self._create([path.name for path in self.ptds]) + + def _create(self, ptd_names: list[str]) -> dict[str, object]: + return create_plain_manifest( + self.root, + { + "pte": Path(self.pte.name), + "source": Path(self.source_receipt.name), + }, + [Path(name) for name in ptd_names], + ) + + def test_plain_requires_exactly_three_ordered_ptds(self) -> None: + for names in ( + [path.name for path in self.ptds[:2]], + [path.name for path in self.ptds] + [self.pte.name], + ): + with self.subTest(count=len(names)): + with self.assertRaisesRegex(ValueError, "exactly three ordered PTDs"): + self._create(names) + + truncated = copy.deepcopy(self.manifest) + dropped = self.ptds[2].name + truncated["artifacts"] = [ + entry for entry in truncated["artifacts"] if entry["path"] != dropped + ] + truncated["ptd_order"] = [ + path for path in truncated["ptd_order"] if path != dropped + ] + self.ptds[2].unlink() + with self.assertRaisesRegex(ValueError, "exactly three ordered PTDs"): + validate_plain_manifest(self.root, truncated) + + def test_plain_ptd_reordering_is_rejected(self) -> None: + reordered = copy.deepcopy(self.manifest) + order = reordered["ptd_order"] + order[0], order[2] = order[2], order[0] + with self.assertRaisesRegex(ValueError, "PTD order does not match"): + validate_plain_manifest(self.root, reordered) + + def test_plain_export_contract_identity(self) -> None: + self.assertEqual(self.manifest["export"], EXPORT_CONTRACT) + mutated = copy.deepcopy(self.manifest) + mutated["export"]["max_seq_len"] = 8961 + with self.assertRaisesRegex(ValueError, "export contract mismatch"): + validate_plain_manifest(self.root, mutated) + + def test_plain_acquisition_identity(self) -> None: + self.assertEqual(self.manifest["acquisition"], CHECKPOINT_ACQUISITION) + self.assertEqual( + self.manifest["model"]["architecture"], ARCHITECTURE_FINGERPRINT + ) + mutated = copy.deepcopy(self.manifest) + mutated["acquisition"]["revision"] = "0" * 40 + with self.assertRaisesRegex(ValueError, "acquisition identity mismatch"): + validate_plain_manifest(self.root, mutated) + + +class CommittedArtifactHygieneTest(unittest.TestCase): + def test_committed_manifests_reference_path_and_hash_only(self) -> None: + documents = sorted((_package_root() / "manifests").glob("*.json")) + self.assertIn("gemma4_e2b_webgpu.json", [path.name for path in documents]) + for path in documents: + with self.subTest(manifest=path.name): + document = json.loads(path.read_text(encoding="utf-8")) + artifacts = document["artifacts"] + self.assertEqual(document["schema_version"], 1) + self.assertNotEqual(artifacts, []) + for entry in artifacts: + self.assertEqual(sorted(entry), _ARTIFACT_KEYS) + self.assertEqual(len(Path(str(entry["path"])).parts), 1) + self.assertIsInstance(entry["bytes"], int) + self.assertRegex(str(entry["sha256"]), "^[0-9a-f]{64}$") + self.assertEqual( + document["ptd_order"], + [entry["path"] for entry in artifacts if entry["role"] == "ptd"], + ) + + def test_committed_plain_manifest_matches_the_export_identity(self) -> None: + document = json.loads( + (_package_root() / "manifests" / "gemma4_e2b_webgpu.json").read_text( + encoding="utf-8" + ) + ) + self.assertEqual(document["export"], EXPORT_CONTRACT) + self.assertEqual(document["acquisition"], CHECKPOINT_ACQUISITION) + self.assertEqual(document["model"]["architecture"], ARCHITECTURE_FINGERPRINT) + self.assertEqual(len(document["ptd_order"]), 3) + self.assertEqual( + sorted(entry["role"] for entry in document["artifacts"]), + ["ptd", "ptd", "ptd", "pte"], + ) + + def test_no_model_binaries_are_committed(self) -> None: + root = _package_root() + self.assertTrue((root / "manifests" / "gemma4_e2b_webgpu.json").is_file()) + self.assertTrue((root / "config" / "e2b_config.json").is_file()) + self.assertEqual( + sorted( + str(path.relative_to(root)) + for path in root.rglob("*") + if path.is_file() and path.suffix in _BINARY_SUFFIXES + ), + [], + ) + + def test_no_internal_paths_leak(self) -> None: + documents = { + "gemma4_webgpu_artifact_manifest.py": Path(gemma4_manifest.__file__), + "backend_webgpu_artifact_manifest.py": Path(backend_manifest.__file__), + } + for path in sorted((_package_root() / "manifests").glob("*.json")): + documents[path.name] = path + self.assertGreaterEqual(len(documents), 3) + for name, path in documents.items(): + text = path.read_text(encoding="utf-8") + for pattern in _internal_patterns(): + with self.subTest(document=name, pattern=pattern.pattern): + self.assertIsNone(pattern.search(text)) diff --git a/examples/models/gemma4/tests/test_webgpu_spec_contract.py b/examples/models/gemma4/tests/test_webgpu_spec_contract.py new file mode 100644 index 00000000000..408de8d9502 --- /dev/null +++ b/examples/models/gemma4/tests/test_webgpu_spec_contract.py @@ -0,0 +1,674 @@ +# 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. + +# pyre-strict + +from __future__ import annotations + +import importlib.util +import os +import re +import subprocess +import tempfile +import unittest + +from pathlib import Path +from typing import Mapping, Sequence + + +GEMMA4_ANCHOR = "examples/models/gemma4/targets.bzl" +SOURCE_ROOT_ENV = "EXECUTORCH_SOURCE_ROOT" +SPEC_RUNNER_HEADER = "examples/models/gemma4/runner/gemma4_spec_runner.h" +SPEC_RUNNER_SOURCE = "examples/models/gemma4/runner/gemma4_spec_runner.cpp" +SPEC_WASM_SOURCE = "examples/models/gemma4/runner/gemma4_spec_wasm.cpp" +GEMMA4_RUNNER_HEADER = "examples/models/gemma4/runner/gemma4_runner.h" +GEMMA4_RUNNER_SOURCE = "examples/models/gemma4/runner/gemma4_runner.cpp" +GEMMA4_TARGETS = "examples/models/gemma4/targets.bzl" +GEMMA4_CMAKE = "examples/models/gemma4/CMakeLists.txt" +GEMMA4_README = "examples/models/gemma4/README.md" +WEBGPU_CMAKE = "backends/webgpu/CMakeLists.txt" +WASM_FACTORY_CONTRACT = "backends/webgpu/scripts/test_gemma4_wasm_factory_contract.sh" +WEBGPU_BACKEND_SOURCE = "backends/webgpu/runtime/WebGPUBackend.cpp" +WEBGPU_EXECUTION_OPTIONS_SOURCE = "backends/webgpu/runtime/WebGPUExecutionOptions.cpp" + +EXPECTED_WASM_EXPORTS = ( + "et_init", + "et_load", + "et_unload", + "et_reset", + "et_prefill_batch", + "et_prefill_step", + "et_step", + "et_mtp_execute_count", + "et_mtp_accepted_drafts", + "et_mtp_buffered_tokens", + "et_mtp_execute", + "et_mtp_execution_attestation", + "et_profile_enable", + "et_profile", +) +EXPECTED_EXPORTED_FUNCTIONS: tuple[str, ...] = tuple( + f"_{symbol}" for symbol in EXPECTED_WASM_EXPORTS +) + ("_malloc", "_free") +EXPECTED_RUNTIME_METHODS = ("ccall", "cwrap", "FS", "HEAP32") + +REQUIRED_RUNNER_SYMBOLS = ( + "Gemma4SpecRunner::accepted_drafts", + "Gemma4SpecRunner::buffered_tokens", + "Gemma4SpecRunner::execute", + "Gemma4SpecRunner::execute_count", + "Gemma4SpecRunner::generate", + "Gemma4SpecRunner::is_loaded", + "Gemma4SpecRunner::load", + "Gemma4SpecRunner::prefill", + "Gemma4SpecRunner::prefill_step", + "Gemma4SpecRunner::profile_json", + "Gemma4SpecRunner::reset", + "Gemma4SpecRunner::set_profiling_enabled", + "Gemma4SpecRunner::step", + "Gemma4SpecRunner::unload", +) + +XNNPACK_SYMBOLS = ( + "weight_cache_option_key", + "workspace_sharing_mode_option_key", + "xnnpack_backend_key", +) + +PUBLIC_GEMMA4_RUNNER_API = ( + "Gemma4Runner", + "load", + "is_loaded", + "generate", + "generate", + "generate_text", + "generate_text", + "generate_vision", + "generate_vision", + "reset", +) + +RESET_LANDMARKS = ( + "impl_->arm_profile();", + "impl_->clear_controller_state();", + "impl_->method_fresh = false;", + "impl_->method_fresh = impl_->method_healthy;", + "impl_->method_healthy = false;", + "->load_method(", + "->unload_method(", +) +EXPECTED_RESET_ORDER = ( + "impl_->method_healthy = false;", + "impl_->method_fresh = false;", + "impl_->clear_controller_state();", + "impl_->arm_profile();", + "impl_->clear_controller_state();", + "->unload_method(", + "impl_->method_healthy = false;", + "impl_->method_fresh = false;", + "->load_method(", + "impl_->method_fresh = impl_->method_healthy;", +) +EXPECTED_HEALTH_LATCH: Mapping[str, tuple[str, ...]] = { + "execute": ("false",) * 10, + "generate": ("false",), + "load": ("false", "false", "false", "true"), + "reset": ("false", "false"), + "step": ("false",), + "unload": ("false",), +} + +_WASM_EXPORT_PATTERN: re.Pattern[str] = re.compile( + r"^ET_WASM_EXPORT\s+[A-Za-z_][\w:*&<>\s]*?\b(et_[a-z0-9_]+)\s*\(", re.M | re.S +) +_DEFINITION_PATTERN: re.Pattern[str] = re.compile( + r"^[A-Za-z_][^\n;{}]*\bGemma4SpecRunner::(\w+)\(", re.M +) +_HEALTH_PATTERN: re.Pattern[str] = re.compile(r"impl_->method_healthy = ([^;]+);") +_BLOCK_COMMENT_PATTERN: re.Pattern[str] = re.compile(r"/\*.*?\*/", re.S) + + +def _root_candidates() -> list[tuple[str, Path | None]]: + override = os.environ.get(SOURCE_ROOT_ENV) + try: + package = importlib.util.find_spec("executorch") + except (ImportError, ValueError): + package = None + staged = list(package.submodule_search_locations or ()) if package else [] + here = Path(__file__).resolve() + walked = next( + (parent for parent in here.parents if (parent / GEMMA4_ANCHOR).is_file()), None + ) + return [ + (f"${SOURCE_ROOT_ENV}", Path(override) if override else None), + ("`executorch` package runfile", Path(staged[0]) if staged else None), + (f"__file__ walk above {here}", walked), + ] + + +def _source_root() -> Path: + attempted: list[str] = [] + for strategy, candidate in _root_candidates(): + attempted.append(f"{strategy} -> {candidate}") + if candidate is not None and (candidate / GEMMA4_ANCHOR).is_file(): + return candidate + raise FileNotFoundError( + f"no ExecuTorch source root containing {GEMMA4_ANCHOR}; " + f"tried {'; '.join(attempted)}" + ) + + +def _read(relative: str) -> str: + path = _source_root() / relative + if not path.is_file(): + raise FileNotFoundError(f"missing source under test: {path}") + return path.read_text(encoding="utf-8") + + +def _verify_product( + source: str, expected_factory: str, expected_output_stem: str +) -> subprocess.CompletedProcess[str]: + with tempfile.TemporaryDirectory() as directory: + javascript = Path(directory) / "product.js" + javascript.write_text(source, encoding="utf-8") + return subprocess.run( + [ + "bash", + str(_source_root() / WASM_FACTORY_CONTRACT), + "--verify-product", + str(javascript), + expected_factory, + expected_output_stem, + ], + check=False, + capture_output=True, + text=True, + ) + + +def wasm_exports(source: str) -> list[str]: + return _WASM_EXPORT_PATTERN.findall(source) + + +def cmake_exported_functions(cmake: str) -> list[str]: + match = re.search(r"-sEXPORTED_FUNCTIONS=\[([^\]]*)\]", cmake) + if match is None: + raise AssertionError("gemma4_spec_browser declares no -sEXPORTED_FUNCTIONS") + return re.findall(r"'([^']+)'", match.group(1)) + + +def cmake_runtime_methods(cmake: str) -> list[str]: + match = re.search(r"-sEXPORTED_RUNTIME_METHODS=\[([^\]]*)\]", cmake) + if match is None: + raise AssertionError( + "gemma4_spec_browser declares no -sEXPORTED_RUNTIME_METHODS" + ) + return re.findall(r"'([^']+)'", match.group(1)) + + +def required_symbol_census(sources: Mapping[str, str]) -> set[str]: + census: set[str] = set() + for text in sources.values(): + census.update(wasm_exports(text)) + census.update( + f"Gemma4SpecRunner::{name}" for name in _DEFINITION_PATTERN.findall(text) + ) + return census + + +def missing_required_symbols(sources: Mapping[str, str]) -> set[str]: + required = set(EXPECTED_WASM_EXPORTS) | set(REQUIRED_RUNNER_SYMBOLS) + return required - required_symbol_census(sources) + + +def _brace_body(text: str, start: int) -> str: + opening = text.index("{", start) + depth = 0 + for index in range(opening, len(text)): + if text[index] == "{": + depth += 1 + elif text[index] == "}": + depth -= 1 + if depth == 0: + return text[opening : index + 1] + raise AssertionError(f"unbalanced braces after offset {start}") + + +def definition_bodies(source: str) -> dict[str, str]: + return { + match.group(1): _brace_body(source, match.end()) + for match in _DEFINITION_PATTERN.finditer(source) + } + + +def wasm_definition_bodies(source: str) -> dict[str, str]: + return { + match.group(1): _brace_body(source, match.end()) + for match in _WASM_EXPORT_PATTERN.finditer(source) + } + + +def ordered_landmarks(body: str, landmarks: Sequence[str]) -> list[str]: + hits: list[tuple[int, str]] = [] + for landmark in landmarks: + start = 0 + while True: + index = body.find(landmark, start) + if index < 0: + break + hits.append((index, landmark)) + start = index + 1 + return [landmark for _, landmark in sorted(hits)] + + +def bzl_rule(text: str, name: str) -> str: + start = text.index(f'name = "{name}",') + return text[start : text.index("\n )\n", start)] + + +def public_api(header: str) -> list[str]: + stripped = _BLOCK_COMMENT_PATTERN.sub("", header) + section = stripped[stripped.index(" public:") : stripped.index(" private:")] + return [ + match.group(1) + for match in re.finditer( + r"^\s{2}(?:[A-Za-z_][\w:<>,\s*&]*?\s)?([A-Za-z_]\w*)\s*\(", + section, + re.M, + ) + ] + + +class BrowserAbiContractTest(unittest.TestCase): + def test_adapter_defines_exactly_the_reviewed_export_list(self) -> None: + self.assertEqual( + wasm_exports(_read(SPEC_WASM_SOURCE)), list(EXPECTED_WASM_EXPORTS) + ) + + def test_cmake_exports_the_adapter_symbols_plus_the_allocator(self) -> None: + exported = cmake_exported_functions(_read(WEBGPU_CMAKE)) + self.assertEqual(exported, list(EXPECTED_EXPORTED_FUNCTIONS)) + self.assertEqual(len(exported), 16) + self.assertEqual( + [symbol[1:] for symbol in exported if symbol.startswith("_et_")], + list(EXPECTED_WASM_EXPORTS), + ) + + def test_cmake_exports_the_reviewed_browser_runtime_methods(self) -> None: + self.assertEqual( + cmake_runtime_methods(_read(WEBGPU_CMAKE)), + list(EXPECTED_RUNTIME_METHODS), + ) + + def test_link_mutant_omitting_the_adapter_fails_the_census(self) -> None: + complete = { + SPEC_WASM_SOURCE: _read(SPEC_WASM_SOURCE), + SPEC_RUNNER_SOURCE: _read(SPEC_RUNNER_SOURCE), + } + self.assertEqual(missing_required_symbols(complete), set()) + without_adapter = {SPEC_RUNNER_SOURCE: complete[SPEC_RUNNER_SOURCE]} + self.assertEqual( + missing_required_symbols(without_adapter), set(EXPECTED_WASM_EXPORTS) + ) + without_runner = {SPEC_WASM_SOURCE: complete[SPEC_WASM_SOURCE]} + self.assertEqual( + missing_required_symbols(without_runner), set(REQUIRED_RUNNER_SYMBOLS) + ) + self.assertEqual( + missing_required_symbols({}), + set(EXPECTED_WASM_EXPORTS) | set(REQUIRED_RUNNER_SYMBOLS), + ) + + def test_adapter_pins_the_method_name_and_tensor_data_path_count(self) -> None: + source = _read(SPEC_WASM_SOURCE) + self.assertIn('constexpr const char* kMethodName = "k2_round";', source) + self.assertIn("constexpr size_t kExpectedTensorDataPaths = 3;", source) + self.assertIn("if (paths.size() != kExpectedTensorDataPaths)", source) + + def test_execution_attestation_reads_the_backend_on_every_call(self) -> None: + bodies = wasm_definition_bodies(_read(SPEC_WASM_SOURCE)) + self.assertIn("et_mtp_execution_attestation", bodies) + body = bodies.get("et_mtp_execution_attestation", "") + self.assertEqual(body.count("webgpu_backend_execution_attestation_json()"), 1) + self.assertIn( + "execution_attestation_json = webgpu_backend_execution_attestation_json();", + body, + ) + self.assertIn("return execution_attestation_json.c_str();", body) + + def test_attestation_reports_observed_pass_and_submit_counts(self) -> None: + backend = _read(WEBGPU_BACKEND_SOURCE) + serializer = _read(WEBGPU_EXECUTION_OPTIONS_SOURCE) + self.assertIn("last_execution_graph->execution_attestation_json()", backend) + self.assertIn('\\"encodedComputePasses\\":', serializer) + self.assertIn('\\"queueSubmitCount\\":', serializer) + + +class SpecRunnerSourceContractTest(unittest.TestCase): + def test_controller_progression_and_bonus_seeding_are_pinned(self) -> None: + header = _read(SPEC_RUNNER_HEADER) + self.assertIn( + "decision.next_position = start_position + output.match_count + 1;", + header, + ) + self.assertIn("decision.next_seed = output.bonus;", header) + self.assertIn("decision.selected.push_back(output.bonus);", header) + + def test_self_consistency_guard_is_pinned(self) -> None: + header = _read(SPEC_RUNNER_HEADER) + self.assertIn( + "if (output.match_count != expected_matches || " + "!valid_token(output.bonus) ||", + header, + ) + self.assertIn( + "output.bonus != output.target_greedy[output.match_count]) {", header + ) + self.assertIn( + "if (start_position < 2 || token_budget == 0 || vocab_size <= 0 ||", + header, + ) + self.assertIn("output.match_count < 0 || output.match_count > 2 ||", header) + self.assertIn("!std::isfinite(output.state_probe)) {", header) + + def test_config_defaults_match_the_export_contract(self) -> None: + header = _read(SPEC_RUNNER_HEADER) + for default in ( + "int64_t vocab_size = 262144;", + "int64_t max_input_length = 512;", + "int64_t target_capacity = 8960;", + "int64_t donor_capacity = 8960;", + 'std::string method_name = "k2_round";', + "int64_t vocab_size = 262144) {", + ): + with self.subTest(default=default): + self.assertIn(default, header) + + def test_load_pins_the_four_input_five_output_vulkan_abi(self) -> None: + source = _read(SPEC_RUNNER_SOURCE) + self.assertIn("meta.num_inputs() != 4 || meta.num_outputs() != 5 ||", source) + self.assertIn( + "meta.num_backends() != 1 || meta.num_instructions() != 1 ||", + source, + ) + self.assertIn('std::string_view(backend.get()) != "VulkanBackend"', source) + self.assertIn("methods->size() != 1 ||", source) + + def test_round_execute_requires_three_rows_and_a_start_aligned_donor( + self, + ) -> None: + source = _read(SPEC_RUNNER_SOURCE) + self.assertIn("(is_round && input_ids.size() != 3)", source) + self.assertIn( + "if ((is_round && (donor_length != start_position || donor_length < 2))", + source, + ) + self.assertIn("if (execution->size() != 5) {", source) + + def test_execute_cannot_request_uncertified_single_compute_pass(self) -> None: + source = _read(SPEC_RUNNER_SOURCE) + body = definition_bodies(source)["execute"] + for forbidden in ( + "WebGPUExecutionOptions", + "single_compute_pass", + "with_webgpu_execution_options(", + ): + with self.subTest(forbidden=forbidden): + self.assertNotIn(forbidden, source) + self.assertIn( + "auto execution = impl_->module->execute(\n" + " impl_->config.method_name,", + body, + ) + + def test_reset_makes_the_next_attestation_fresh(self) -> None: + reset = definition_bodies(_read(SPEC_RUNNER_SOURCE))["reset"] + unload = reset.index("->unload_method(") + reload = reset.index("->load_method(") + self.assertLess(unload, reload) + + backend = _read(WEBGPU_BACKEND_SOURCE) + destroy_start = backend.index("void WebGPUBackend::destroy(") + destroy = _brace_body(backend, destroy_start) + self.assertIn("if (last_execution_graph == graph)", destroy) + self.assertIn("last_execution_graph = nullptr;", destroy) + + bodies = wasm_definition_bodies(_read(SPEC_WASM_SOURCE)) + self.assertIn("et_mtp_execution_attestation", bodies) + attestation = bodies.get("et_mtp_execution_attestation", "") + self.assertIn("webgpu_backend_execution_attestation_json()", attestation) + + def test_reset_clears_state_then_unloads_then_reloads(self) -> None: + body = definition_bodies(_read(SPEC_RUNNER_SOURCE))["reset"] + self.assertEqual( + tuple(ordered_landmarks(body, RESET_LANDMARKS)), EXPECTED_RESET_ORDER + ) + self.assertIn( + "error == Error::Ok && verify_context(impl_->context.get())", body + ) + + def test_method_healthy_latch_is_cleared_only_by_reset_and_load(self) -> None: + source = _read(SPEC_RUNNER_SOURCE) + bodies = definition_bodies(source) + observed = { + name: tuple(_HEALTH_PATTERN.findall(body)) + for name, body in bodies.items() + if _HEALTH_PATTERN.search(body) + } + self.assertEqual(observed, EXPECTED_HEALTH_LATCH) + self.assertEqual(len(_HEALTH_PATTERN.findall(source)), 19) + + def test_unload_surfaces_a_failed_method_unload(self) -> None: + body = definition_bodies(_read(SPEC_RUNNER_SOURCE))["unload"] + self.assertIn("return method_unloaded ? Error::Ok : Error::Internal;", body) + self.assertIn("impl_->clear_controller_state();", body) + self.assertIn("destroy_webgpu_context(*impl_->context);", body) + + def test_is_loaded_requires_a_module_and_a_healthy_method(self) -> None: + body = definition_bodies(_read(SPEC_RUNNER_SOURCE))["is_loaded"] + self.assertIn( + "return impl_->module != nullptr && impl_->method_healthy &&", body + ) + self.assertIn( + "impl_->context != nullptr && verify_context(impl_->context.get())", body + ) + + +class XnnpackPreservationTest(unittest.TestCase): + def test_spec_sources_reference_no_xnnpack_symbol(self) -> None: + control = _read(GEMMA4_RUNNER_SOURCE) + for symbol in XNNPACK_SYMBOLS: + with self.subTest(symbol=symbol): + self.assertIn(symbol, control) + for relative in (SPEC_RUNNER_HEADER, SPEC_RUNNER_SOURCE, SPEC_WASM_SOURCE): + text = _read(relative) + for symbol in XNNPACK_SYMBOLS + ("xnnpack", "XNNPACK"): + with self.subTest(source=relative, symbol=symbol): + self.assertNotIn(symbol, text) + + def test_spec_targets_declare_no_xnnpack_dependency(self) -> None: + targets = _read(GEMMA4_TARGETS) + for name in ("gemma4_spec_runner", "gemma4_spec_wasm_adapter"): + with self.subTest(target=name): + self.assertNotIn("xnnpack", bzl_rule(targets, name)) + + def test_public_gemma4_runner_target_is_unchanged(self) -> None: + rule = bzl_rule(_read(GEMMA4_TARGETS), "gemma4_runner") + self.assertIn('srcs = [\n "runner/gemma4_runner.cpp",\n', rule) + for header in ( + "runner/gemma4_runner.h", + "runner/gemma4_stats.h", + "runner/generation_config.h", + ): + with self.subTest(header=header): + self.assertIn(f'"{header}",', rule) + self.assertIn('"//executorch/backends/xnnpack:xnnpack_interface",', rule) + self.assertIn("deps = _KERNEL_BACKEND_DEPS + [", rule) + + def test_public_gemma4_runner_api_is_unchanged(self) -> None: + self.assertEqual( + public_api(_read(GEMMA4_RUNNER_HEADER)), list(PUBLIC_GEMMA4_RUNNER_API) + ) + + +class SpecBuildContractTest(unittest.TestCase): + def test_spec_browser_compiles_the_adapter_runner_and_links_the_loader( + self, + ) -> None: + cmake = _read(WEBGPU_CMAKE) + start = cmake.index("add_executable(\n gemma4_spec_browser") + block = cmake[start : cmake.index("-sEXPORTED_FUNCTIONS", start)] + for fragment in ( + "examples/models/gemma4/runner/gemma4_spec_runner.cpp", + "examples/models/gemma4/runner/gemma4_spec_wasm.cpp", + "webgpu_backend webgpu_model_loader", + "extension_tensor", + "--use-port=emdawnwebgpu", + ): + with self.subTest(fragment=fragment): + self.assertIn(fragment, block) + + def test_mtp_target_consumes_validated_factory_and_output_cache_strings( + self, + ) -> None: + cmake = _read(WEBGPU_CMAKE) + plain_start = cmake.index("add_executable(\n gemma4_plain_wasm") + spec_start = cmake.index("add_executable(\n gemma4_spec_browser") + plain = cmake[plain_start:spec_start] + spec = cmake[spec_start : cmake.index("\nendif()", spec_start)] + self.assertIn( + "set(GEMMA4_SPEC_WASM_EXPORT_NAME\n" + ' "createGemma4Mtp"\n' + " CACHE STRING", + cmake, + ) + self.assertIn( + "set(GEMMA4_SPEC_WASM_OUTPUT_NAME\n" + ' "gemma4_mtp"\n' + " CACHE STRING", + cmake, + ) + self.assertIn("include(cmake/ValidateGemma4WasmNames.cmake)", cmake) + self.assertIn( + "validate_gemma4_wasm_names(\n" + " GEMMA4_SPEC_WASM_EXPORT_NAME GEMMA4_SPEC_WASM_OUTPUT_NAME\n" + " )", + cmake, + ) + self.assertIn("-sEXPORT_NAME=${GEMMA4_SPEC_WASM_EXPORT_NAME}", spec) + self.assertIn('OUTPUT_NAME "${GEMMA4_SPEC_WASM_OUTPUT_NAME}"', spec) + self.assertNotIn("GEMMA4_SPEC_WASM_EXPORT_NAME", plain) + self.assertNotIn("GEMMA4_SPEC_WASM_OUTPUT_NAME", plain) + self.assertIn("-sEXPORT_NAME=createWebGPULlama", plain) + self.assertIn('OUTPUT_NAME "webgpu_llama"', plain) + + def test_readme_builds_distinct_wall_and_profile_products(self) -> None: + readme = _read(GEMMA4_README) + wall_start = readme.index('emcmake cmake "${COMMON[@]}" -B "$WALL_BUILD"') + profile_start = readme.index('emcmake cmake "${COMMON[@]}" -B "$PROFILE_BUILD"') + wall = readme[wall_start:profile_start] + profile = readme[profile_start : readme.index("\n```", profile_start)] + for block, required, forbidden in ( + ( + wall, + ( + "-DEXECUTORCH_BUILD_WEBGPU_PROFILING=OFF", + "-DGEMMA4_SPEC_WASM_EXPORT_NAME=createGemma4Mtp", + "-DGEMMA4_SPEC_WASM_OUTPUT_NAME=gemma4_mtp", + ), + ("PROFILING=ON", "createGemma4MtpProfile", "gemma4_mtp_profile"), + ), + ( + profile, + ( + "-DEXECUTORCH_BUILD_WEBGPU_PROFILING=ON", + "-DGEMMA4_SPEC_WASM_EXPORT_NAME=createGemma4MtpProfile", + "-DGEMMA4_SPEC_WASM_OUTPUT_NAME=gemma4_mtp_profile", + ), + ("PROFILING=OFF", "=createGemma4Mtp\n", "=gemma4_mtp\n"), + ), + ): + for fragment in required: + with self.subTest(fragment=fragment): + self.assertIn(fragment, block) + for fragment in forbidden: + with self.subTest(forbidden=fragment): + self.assertNotIn(fragment, block) + self.assertIn("--target gemma4_plain_wasm gemma4_spec_browser", block) + + for fragment in ( + "plain-profile-recipe.json", + "browser_gemma4_mtp/gemma4_mtp_profile.js", + "browser_gemma4_mtp/gemma4_mtp_profile.wasm", + "--plain-profile-javascript", + "--plain-profile-wasm", + "--plain-profile-recipe", + ): + with self.subTest(fragment=fragment): + self.assertIn(fragment, readme) + + def test_product_verifier_binds_factory_and_requested_wasm(self) -> None: + accepted = _verify_product( + "var createGemma4Mtp = async function(options) {" + "options.locateFile('gemma4_mtp.wasm');};", + "createGemma4Mtp", + "gemma4_mtp", + ) + self.assertEqual(accepted.returncode, 0, accepted.stderr) + + mutations = { + "wrong factory": ( + "var createGemma4MtpProfile = async function(options) {" + "options.locateFile('gemma4_mtp.wasm');};", + "createGemma4Mtp", + "gemma4_mtp", + ), + "wrong wasm": ( + "var createGemma4Mtp = async function(options) {" + "options.locateFile('wrong.wasm');};", + "createGemma4Mtp", + "gemma4_mtp", + ), + "extra factory": ( + "var createGemma4Mtp = async function(options) {" + "options.locateFile('gemma4_mtp.wasm');};" + "var createWebGPULlama = function() {};", + "createGemma4Mtp", + "gemma4_mtp", + ), + } + for label, (source, factory, output_stem) in mutations.items(): + with self.subTest(label=label): + rejected = _verify_product(source, factory, output_stem) + self.assertNotEqual(rejected.returncode, 0) + + def test_spec_browser_has_the_full_production_link_closure(self) -> None: + cmake = _read(WEBGPU_CMAKE) + start = cmake.index("add_executable(\n gemma4_spec_browser") + block = cmake[start : cmake.index("\nendif()", start)] + for fragment in ( + "-fexceptions", + '"--use-port=emdawnwebgpu"', + '"-sASYNCIFY"', + '"-sALLOW_MEMORY_GROWTH=1"', + '"-sMAXIMUM_MEMORY=4GB"', + '"-sFORCE_FILESYSTEM=1"', + '"--no-entry"', + "'HEAP32'", + '"-sSTACK_SIZE=8388608"', + '"-sASYNCIFY_STACK_SIZE=1048576"', + '"-sMODULARIZE=1"', + ): + with self.subTest(fragment=fragment): + self.assertIn(fragment, block) + self.assertNotIn("-sNO_ENTRY", block) + + def test_native_spec_runner_is_guarded_by_the_loader_target(self) -> None: + cmake = _read(GEMMA4_CMAKE) + self.assertIn("if(TARGET webgpu_backend AND TARGET webgpu_model_loader)", cmake) + self.assertIn( + "add_library(gemma4_spec_runner runner/gemma4_spec_runner.cpp)", cmake + ) diff --git a/examples/models/gemma4/webgpu_artifact_manifest.py b/examples/models/gemma4/webgpu_artifact_manifest.py index 7bc73fab6cf..7ddc04f1b14 100644 --- a/examples/models/gemma4/webgpu_artifact_manifest.py +++ b/examples/models/gemma4/webgpu_artifact_manifest.py @@ -156,6 +156,7 @@ "[ExecuTorch][WebGPU] Add Gemma 4 MTP operator and route support", "[ExecuTorch][WebGPU] Add Gemma 4 MTP export path", "[ExecuTorch][WebGPU] Add Gemma 4 speculative decode runtime", + "[ExecuTorch][WebGPU] Add Gemma 4 MTP and speculative-decode source-closure tests", ) MTP_EXPORT_CONTRACT: dict[str, object] = { "assistant_calls_per_round": 2,