From 85462467dcc36d47c17ff85991e5248ff76e2b35 Mon Sep 17 00:00:00 2001 From: Julian Ng-Thow-Hing Date: Fri, 7 Aug 2026 11:57:08 -0700 Subject: [PATCH] Update [ghstack-poisoned] --- backends/webgpu/CMakeLists.txt | 58 +- .../cmake/ValidateGemma4WasmNames.cmake | 35 + .../test_gemma4_wasm_factory_contract.sh | 155 ++ examples/models/gemma4/CMakeLists.txt | 21 + examples/models/gemma4/README.md | 185 +- examples/models/gemma4/export_speculative.py | 66 +- .../models/gemma4/runner/gemma4_spec_main.cpp | 156 ++ .../gemma4/runner/gemma4_spec_runner.cpp | 1003 ++++++++ .../models/gemma4/runner/gemma4_spec_runner.h | 185 ++ .../models/gemma4/runner/gemma4_spec_wasm.cpp | 194 ++ examples/models/gemma4/targets.bzl | 37 + .../models/gemma4/webgpu_artifact_manifest.py | 2202 ++++++++++++++++- 12 files changed, 4221 insertions(+), 76 deletions(-) create mode 100644 backends/webgpu/cmake/ValidateGemma4WasmNames.cmake create mode 100755 backends/webgpu/scripts/test_gemma4_wasm_factory_contract.sh create mode 100644 examples/models/gemma4/runner/gemma4_spec_main.cpp create mode 100644 examples/models/gemma4/runner/gemma4_spec_runner.cpp create mode 100644 examples/models/gemma4/runner/gemma4_spec_runner.h create mode 100644 examples/models/gemma4/runner/gemma4_spec_wasm.cpp diff --git a/backends/webgpu/CMakeLists.txt b/backends/webgpu/CMakeLists.txt index a112f30d275..919a849a7a6 100644 --- a/backends/webgpu/CMakeLists.txt +++ b/backends/webgpu/CMakeLists.txt @@ -117,6 +117,19 @@ executorch_target_link_options_shared_lib(webgpu_backend) set_property(TARGET webgpu_backend PROPERTY CXX_STANDARD 17) if(EMSCRIPTEN) + include(cmake/ValidateGemma4WasmNames.cmake) + set(GEMMA4_SPEC_WASM_EXPORT_NAME + "createGemma4Mtp" + CACHE STRING "JavaScript factory exported by the Gemma 4 MTP WASM module" + ) + set(GEMMA4_SPEC_WASM_OUTPUT_NAME + "gemma4_mtp" + CACHE STRING "Output file stem for the Gemma 4 MTP WASM module" + ) + validate_gemma4_wasm_names( + GEMMA4_SPEC_WASM_EXPORT_NAME GEMMA4_SPEC_WASM_OUTPUT_NAME + ) + add_executable( gemma4_plain_wasm ${EXECUTORCH_ROOT}/examples/models/gemma4/runner/gemma4_plain_wasm.cpp @@ -125,7 +138,8 @@ if(EMSCRIPTEN) gemma4_plain_wasm PRIVATE $ ) target_link_libraries( - gemma4_plain_wasm PRIVATE webgpu_backend webgpu_model_loader extension_tensor + gemma4_plain_wasm PRIVATE webgpu_backend webgpu_model_loader + extension_tensor ) target_compile_options( gemma4_plain_wasm PRIVATE -fexceptions "--use-port=emdawnwebgpu" @@ -159,6 +173,48 @@ if(EMSCRIPTEN) "${CMAKE_CURRENT_BINARY_DIR}/browser_gemma4_plain" CXX_STANDARD 17 ) + add_executable( + gemma4_spec_browser + ${EXECUTORCH_ROOT}/examples/models/gemma4/runner/gemma4_spec_runner.cpp + ${EXECUTORCH_ROOT}/examples/models/gemma4/runner/gemma4_spec_wasm.cpp + ) + target_include_directories( + gemma4_spec_browser PRIVATE $ + ) + target_link_libraries( + gemma4_spec_browser PRIVATE webgpu_backend webgpu_model_loader + extension_tensor + ) + target_compile_options(gemma4_spec_browser PRIVATE -fexceptions) + if(EXECUTORCH_BUILD_WEBGPU_PROFILING) + target_compile_definitions( + gemma4_spec_browser PRIVATE WGPU_BACKEND_ENABLE_PROFILING + ) + endif() + target_link_options( + gemma4_spec_browser + PRIVATE + -fexceptions + "--use-port=emdawnwebgpu" + "-sASYNCIFY" + "-sALLOW_MEMORY_GROWTH=1" + "-sMAXIMUM_MEMORY=4GB" + "-sFORCE_FILESYSTEM=1" + "--no-entry" + "-sEXPORTED_FUNCTIONS=['_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','_malloc','_free']" + "-sEXPORTED_RUNTIME_METHODS=['ccall','cwrap','FS','HEAP32']" + "-sSTACK_SIZE=8388608" + "-sASYNCIFY_STACK_SIZE=1048576" + "-sMODULARIZE=1" + "-sEXPORT_NAME=${GEMMA4_SPEC_WASM_EXPORT_NAME}" + ) + set_target_properties( + gemma4_spec_browser + PROPERTIES OUTPUT_NAME "${GEMMA4_SPEC_WASM_OUTPUT_NAME}" + RUNTIME_OUTPUT_DIRECTORY + "${CMAKE_CURRENT_BINARY_DIR}/browser_gemma4_mtp" + CXX_STANDARD 17 + ) endif() install( diff --git a/backends/webgpu/cmake/ValidateGemma4WasmNames.cmake b/backends/webgpu/cmake/ValidateGemma4WasmNames.cmake new file mode 100644 index 00000000000..cfa5e1ccfb9 --- /dev/null +++ b/backends/webgpu/cmake/ValidateGemma4WasmNames.cmake @@ -0,0 +1,35 @@ +# 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. + +function(validate_gemma4_wasm_names export_variable output_variable) + if(NOT DEFINED ${export_variable}) + message(FATAL_ERROR "${export_variable} must be defined") + endif() + if(NOT DEFINED ${output_variable}) + message(FATAL_ERROR "${output_variable} must be defined") + endif() + + set(export_name "${${export_variable}}") + set(output_name "${${output_variable}}") + if(NOT export_name MATCHES "^[A-Za-z_$][A-Za-z0-9_$]*$") + message( + FATAL_ERROR + "${export_variable} must be a JavaScript identifier: '${export_name}'" + ) + endif() + if(NOT output_name MATCHES "^[A-Za-z0-9][A-Za-z0-9._-]*$") + message( + FATAL_ERROR + "${output_variable} must be a file-name stem: '${output_name}'" + ) + endif() +endfunction() + +if(CMAKE_SCRIPT_MODE_FILE AND GEMMA4_VALIDATE_WASM_NAMES) + validate_gemma4_wasm_names( + GEMMA4_SPEC_WASM_EXPORT_NAME GEMMA4_SPEC_WASM_OUTPUT_NAME + ) +endif() diff --git a/backends/webgpu/scripts/test_gemma4_wasm_factory_contract.sh b/backends/webgpu/scripts/test_gemma4_wasm_factory_contract.sh new file mode 100755 index 00000000000..d3c59e44bc2 --- /dev/null +++ b/backends/webgpu/scripts/test_gemma4_wasm_factory_contract.sh @@ -0,0 +1,155 @@ +#!/bin/bash +# 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. + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +VALIDATOR="${SCRIPT_DIR}/../cmake/ValidateGemma4WasmNames.cmake" +VALIDATION_ERROR="must be" + +validate_names() { + local export_name="$1" + local output_name="$2" + cmake \ + -DGEMMA4_VALIDATE_WASM_NAMES=ON \ + "-DGEMMA4_SPEC_WASM_EXPORT_NAME:STRING=${export_name}" \ + "-DGEMMA4_SPEC_WASM_OUTPUT_NAME:STRING=${output_name}" \ + -P "${VALIDATOR}" +} + +expect_invalid() { + local export_name="$1" + local output_name="$2" + local expected="$3" + local output + if output="$(validate_names "${export_name}" "${output_name}" 2>&1)"; then + echo "ERROR: invalid Gemma 4 WASM name pair was accepted" >&2 + return 1 + fi + case "${output}" in + *"${expected}"*"${VALIDATION_ERROR}"*) ;; + *) + printf 'ERROR: unexpected CMake validation failure:\n%s\n' "${output}" >&2 + return 1 + ;; + esac +} + +validate_name_matrix() { + validate_names 'create$Gemma4Mtp_1' '1gemma4_mtp-profile.1' + + expect_invalid '' gemma4_mtp GEMMA4_SPEC_WASM_EXPORT_NAME + expect_invalid 'bad;name' gemma4_mtp GEMMA4_SPEC_WASM_EXPORT_NAME + expect_invalid 'bad name' gemma4_mtp GEMMA4_SPEC_WASM_EXPORT_NAME + expect_invalid 'bad/name' gemma4_mtp GEMMA4_SPEC_WASM_EXPORT_NAME + expect_invalid 'bad\name' gemma4_mtp GEMMA4_SPEC_WASM_EXPORT_NAME + expect_invalid . gemma4_mtp GEMMA4_SPEC_WASM_EXPORT_NAME + expect_invalid .. gemma4_mtp GEMMA4_SPEC_WASM_EXPORT_NAME + expect_invalid .hidden gemma4_mtp GEMMA4_SPEC_WASM_EXPORT_NAME + expect_invalid 1factory gemma4_mtp GEMMA4_SPEC_WASM_EXPORT_NAME + expect_invalid bad-name gemma4_mtp GEMMA4_SPEC_WASM_EXPORT_NAME + + expect_invalid createGemma4Mtp '' GEMMA4_SPEC_WASM_OUTPUT_NAME + expect_invalid createGemma4Mtp 'bad;name' GEMMA4_SPEC_WASM_OUTPUT_NAME + expect_invalid createGemma4Mtp 'bad name' GEMMA4_SPEC_WASM_OUTPUT_NAME + expect_invalid createGemma4Mtp 'bad/name' GEMMA4_SPEC_WASM_OUTPUT_NAME + expect_invalid createGemma4Mtp 'bad\name' GEMMA4_SPEC_WASM_OUTPUT_NAME + expect_invalid createGemma4Mtp . GEMMA4_SPEC_WASM_OUTPUT_NAME + expect_invalid createGemma4Mtp .. GEMMA4_SPEC_WASM_OUTPUT_NAME + expect_invalid createGemma4Mtp .hidden GEMMA4_SPEC_WASM_OUTPUT_NAME + + echo "Gemma 4 WASM name validation passed" +} + +verify_product() { + local javascript="$1" + local expected_factory="$2" + local expected_output_stem="$3" + node - "${javascript}" "${expected_factory}" "${expected_output_stem}" <<'NODE' +const fs = require("fs"); +const vm = require("vm"); + +const [javascriptPath, expectedFactory, expectedOutputStem] = process.argv.slice(2); +const knownFactories = [ + "createWebGPULlama", + "createGemma4Mtp", + "createGemma4MtpProfile", +]; + +async function main() { + const context = vm.createContext({}); + for (const commonJsName of ["module", "exports", "require"]) { + if (vm.runInContext(`typeof ${commonJsName}`, context) !== "undefined") { + throw new Error(`fresh VM unexpectedly defines ${commonJsName}`); + } + } + vm.runInContext(fs.readFileSync(javascriptPath, "utf8"), context, { + filename: javascriptPath, + }); + for (const factory of knownFactories) { + const type = vm.runInContext(`typeof ${factory}`, context); + if (factory === expectedFactory) { + if (type !== "function") { + throw new Error(`expected factory ${factory} is not callable`); + } + } else if (type !== "undefined") { + throw new Error(`unexpected Gemma factory published: ${factory}`); + } + } + + const requests = []; + const sentinel = new Error("stop before WASM fetch"); + let rejected = false; + try { + await context[expectedFactory]({ + locateFile(path) { + requests.push(path); + throw sentinel; + }, + }); + } catch (_error) { + rejected = true; + } + if (!rejected) { + throw new Error("modularized factory resolved before the locateFile sentinel"); + } + if (requests.length !== 1) { + throw new Error(`expected one WASM request, observed ${requests.length}`); + } + const expectedWasm = `${expectedOutputStem}.wasm`; + if (requests[0] !== expectedWasm) { + throw new Error(`expected WASM request ${expectedWasm}, observed ${requests[0]}`); + } +} + +main().catch((error) => { + console.error(error.message); + process.exitCode = 1; +}); +NODE +} + +case "${1:-}" in + --validate-names) + if [[ "$#" -ne 1 ]]; then + echo "usage: $0 --validate-names" >&2 + exit 2 + fi + validate_name_matrix + ;; + --verify-product) + if [[ "$#" -ne 4 ]]; then + echo "usage: $0 --verify-product JS EXPECTED_FACTORY EXPECTED_OUTPUT_STEM" >&2 + exit 2 + fi + verify_product "$2" "$3" "$4" + ;; + *) + echo "usage: $0 --validate-names | --verify-product JS EXPECTED_FACTORY EXPECTED_OUTPUT_STEM" >&2 + exit 2 + ;; +esac diff --git a/examples/models/gemma4/CMakeLists.txt b/examples/models/gemma4/CMakeLists.txt index d4e05e70306..455a43ef3f6 100644 --- a/examples/models/gemma4/CMakeLists.txt +++ b/examples/models/gemma4/CMakeLists.txt @@ -82,3 +82,24 @@ target_include_directories( ) target_link_libraries(gemma4_e2e_runner PUBLIC ${link_libraries}) target_compile_options(gemma4_e2e_runner PUBLIC ${_common_compile_options}) + +if(TARGET webgpu_backend AND TARGET webgpu_model_loader) + add_library(gemma4_spec_runner runner/gemma4_spec_runner.cpp) + target_include_directories( + gemma4_spec_runner PUBLIC ${_common_include_directories} + ) + target_link_libraries( + gemma4_spec_runner PUBLIC webgpu_backend webgpu_model_loader + extension_tensor + ) + target_compile_options(gemma4_spec_runner PRIVATE -fexceptions) + if(EXECUTORCH_BUILD_WEBGPU_PROFILING) + target_compile_definitions( + gemma4_spec_runner PRIVATE WGPU_BACKEND_ENABLE_PROFILING + ) + endif() + + add_executable(gemma4_spec_runner_cli runner/gemma4_spec_main.cpp) + target_link_libraries(gemma4_spec_runner_cli PRIVATE gemma4_spec_runner) + target_compile_options(gemma4_spec_runner_cli PRIVATE -fexceptions) +endif() diff --git a/examples/models/gemma4/README.md b/examples/models/gemma4/README.md index 7289978f0c3..cf132e1ebd5 100644 --- a/examples/models/gemma4/README.md +++ b/examples/models/gemma4/README.md @@ -71,8 +71,9 @@ buck2 run fbcode//executorch/examples/models/gemma4:webgpu_artifact_manifest -- validate-acquisition --checkpoint-root /tmp/gemma4-e2b-it ``` -From clean fbsource and ExecuTorch OSS checkouts, seal the reviewed plain-Gemma -source union and generator-derived WGSL closure before exporting the model: +From clean fbsource and ExecuTorch OSS checkouts, seal the reviewed Gemma +production source union and generator-derived WGSL closure before exporting +either model: ```bash : "${FBSOURCE_ROOT:?set the clean fbsource checkout root}" @@ -271,6 +272,186 @@ When formatting the answer, first output the transcription in {SOURCE_LANGUAGE}, | E2B gemma4_vision.pte | 4.3 GB | 798ms | 2.73s | 134 tok/s | 6 tok/s | 3.83s | 10.14s | 1884 MB | 2600 MB | | E4B gemma4_vision.pte | 6.2 GB | 1.36s | 2.44s | 85 tok/s | 4 tok/s | 4.17s | 14.62s | 3232 MB | 3950 MB | +## WebGPU speculative decoding + +The WebGPU MTP path loads one fully delegated `k2_round` method and three +caller-ordered external tensor-data files. The runner does not locate, +download, validate, or copy model artifacts. + +Build the reusable native controller and token-ID CLI with: + +```bash +buck2 build fbcode//executorch/examples/models/gemma4:gemma4_spec_runner_cli +``` + +Acquire the assistant checkpoint at its pinned revision, validate it, and export +the source-verified `k2_round` program: + +```bash +hf download google/gemma-4-E2B-it-qat-q4_0-unquantized-assistant \ + config.json model.safetensors \ + --revision ebc7e1a211354561464cb82ed6d886792138dcb6 \ + --local-dir /tmp/gemma4-e2b-assistant +python -m executorch.examples.models.gemma4.webgpu_artifact_manifest \ + validate-assistant-acquisition \ + --checkpoint-root /tmp/gemma4-e2b-assistant +python -m executorch.examples.models.gemma4.export_speculative \ + --target-checkpoint /tmp/gemma4-e2b-it \ + --assistant-checkpoint /tmp/gemma4-e2b-assistant \ + --output /tmp/gemma4-mtp/model.pte \ + --receipt /tmp/gemma4-mtp.json \ + --source-receipt /tmp/gemma4-source-receipt.json \ + --max-seq-len 8960 --max-input-len 512 +``` + +From a clean OSS checkout with Emscripten 4.0.10 activated, build the wall and +profiling plain/MTP adapter pairs in separate directories with: + +```bash +export EXECUTORCH_ROOT="$PWD" +export WALL_BUILD="$EXECUTORCH_ROOT/cmake-out-gemma4-webgpu-wall" +export PROFILE_BUILD="$EXECUTORCH_ROOT/cmake-out-gemma4-webgpu-profile" +COMMON=( + -S "$EXECUTORCH_ROOT" -GNinja -DCMAKE_BUILD_TYPE=Release + -DPYTHON_EXECUTABLE="$EXECUTORCH_ROOT/.venv/bin/python" + -DEXECUTORCH_BUILD_WEBGPU=ON -DEXECUTORCH_BUILD_WEBGPU_TEST=OFF + -DEXECUTORCH_BUILD_WASM=ON -DEXECUTORCH_BUILD_XNNPACK=OFF + -DEXECUTORCH_BUILD_CPUINFO=ON -DEXECUTORCH_BUILD_PTHREADPOOL=ON + -DEXECUTORCH_BUILD_EXTENSION_MODULE=ON + -DEXECUTORCH_BUILD_EXTENSION_DATA_LOADER=ON + -DEXECUTORCH_BUILD_EXTENSION_TENSOR=ON + -DEXECUTORCH_BUILD_EXTENSION_FLAT_TENSOR=ON + -DEXECUTORCH_BUILD_EXTENSION_NAMED_DATA_MAP=ON +) +emcmake cmake "${COMMON[@]}" -B "$WALL_BUILD" \ + -DEXECUTORCH_BUILD_WEBGPU_PROFILING=OFF \ + -DGEMMA4_SPEC_WASM_EXPORT_NAME=createGemma4Mtp \ + -DGEMMA4_SPEC_WASM_OUTPUT_NAME=gemma4_mtp +cmake --build "$WALL_BUILD" \ + --target gemma4_plain_wasm gemma4_spec_browser -j"$(nproc)" +emcmake cmake "${COMMON[@]}" -B "$PROFILE_BUILD" \ + -DEXECUTORCH_BUILD_WEBGPU_PROFILING=ON \ + -DGEMMA4_SPEC_WASM_EXPORT_NAME=createGemma4MtpProfile \ + -DGEMMA4_SPEC_WASM_OUTPUT_NAME=gemma4_mtp_profile +cmake --build "$PROFILE_BUILD" \ + --target gemma4_plain_wasm gemma4_spec_browser -j"$(nproc)" +``` + +Verify each modularized JavaScript product publishes only its recorded factory +and requests its recorded WASM basename: + +```bash +FACTORY_GATE="$EXECUTORCH_ROOT/backends/webgpu/scripts/test_gemma4_wasm_factory_contract.sh" +bash "$FACTORY_GATE" --verify-product \ + "$WALL_BUILD/backends/webgpu/browser_gemma4_plain/webgpu_llama.js" \ + createWebGPULlama webgpu_llama +bash "$FACTORY_GATE" --verify-product \ + "$PROFILE_BUILD/backends/webgpu/browser_gemma4_plain/webgpu_llama.js" \ + createWebGPULlama webgpu_llama +bash "$FACTORY_GATE" --verify-product \ + "$WALL_BUILD/backends/webgpu/browser_gemma4_mtp/gemma4_mtp.js" \ + createGemma4Mtp gemma4_mtp +bash "$FACTORY_GATE" --verify-product \ + "$PROFILE_BUILD/backends/webgpu/browser_gemma4_mtp/gemma4_mtp_profile.js" \ + createGemma4MtpProfile gemma4_mtp_profile +``` + +Run a manifest-staged PTE and its three ordered PTDs with: + +```bash +gemma4_spec_runner_cli --pte model.pte \ + --ptd part0.ptd --ptd part1.ptd --ptd part2.ptd \ + --prompt-ids 2,123,456 --max-new-tokens 32 +``` + +Write canonical reproduction recipes, then bind independently validated plain +and MTP receipts to the provided wall/profile runtime bytes using: + +```bash +python -m executorch.examples.models.gemma4.webgpu_artifact_manifest \ + create-build-recipe --model plain --flavor wall \ + --output plain-wall-recipe.json +python -m executorch.examples.models.gemma4.webgpu_artifact_manifest \ + create-build-recipe --model plain --flavor profile \ + --output plain-profile-recipe.json +python -m executorch.examples.models.gemma4.webgpu_artifact_manifest \ + create-build-recipe --model mtp --flavor wall \ + --output mtp-wall-recipe.json +python -m executorch.examples.models.gemma4.webgpu_artifact_manifest \ + create-build-recipe --model mtp --flavor profile \ + --output mtp-profile-recipe.json +SOURCE_MANIFEST=/tmp/gemma4-source-manifest.json +WGSL_MANIFEST=/tmp/gemma4-wgsl-manifest.json +PLAIN_ROOT=/tmp/gemma4-webgpu +PLAIN_MANIFEST=/tmp/gemma4-e2b-webgpu.json +MTP_ROOT=/tmp/gemma4-mtp +MTP_MANIFEST=/tmp/gemma4-mtp.json +python -m executorch.examples.models.gemma4.webgpu_artifact_manifest \ + create-runtime-source --output runtime-source.json \ + --fbsource-root "$FBSOURCE_ROOT" --oss-root "$OSS_ROOT" \ + --backend-root "$FBSOURCE_ROOT/xplat/executorch/backends/webgpu" \ + --plain-root "$PLAIN_ROOT" --mtp-root "$MTP_ROOT" \ + --source-manifest "$SOURCE_MANIFEST" --wgsl-manifest "$WGSL_MANIFEST" \ + --plain-manifest "$PLAIN_MANIFEST" --mtp-manifest "$MTP_MANIFEST" \ + --plain-wall-javascript "$WALL_BUILD/backends/webgpu/browser_gemma4_plain/webgpu_llama.js" \ + --plain-wall-wasm "$WALL_BUILD/backends/webgpu/browser_gemma4_plain/webgpu_llama.wasm" \ + --plain-wall-recipe plain-wall-recipe.json \ + --plain-profile-javascript "$PROFILE_BUILD/backends/webgpu/browser_gemma4_plain/webgpu_llama.js" \ + --plain-profile-wasm "$PROFILE_BUILD/backends/webgpu/browser_gemma4_plain/webgpu_llama.wasm" \ + --plain-profile-recipe plain-profile-recipe.json \ + --mtp-wall-javascript "$WALL_BUILD/backends/webgpu/browser_gemma4_mtp/gemma4_mtp.js" \ + --mtp-wall-wasm "$WALL_BUILD/backends/webgpu/browser_gemma4_mtp/gemma4_mtp.wasm" \ + --mtp-wall-recipe mtp-wall-recipe.json \ + --mtp-profile-javascript "$PROFILE_BUILD/backends/webgpu/browser_gemma4_mtp/gemma4_mtp_profile.js" \ + --mtp-profile-wasm "$PROFILE_BUILD/backends/webgpu/browser_gemma4_mtp/gemma4_mtp_profile.wasm" \ + --mtp-profile-recipe mtp-profile-recipe.json +buck2 run fbcode//executorch/examples/models/gemma4:generate_target_prefill_oracle -- \ + --checkpoints /tmp/gemma4-e2b-it \ + --runtime-source-receipt runtime-source.json \ + --contexts 128,511,512,513,514,1024,2048,4096,4097,8192 \ + --output target-prefill.json +python -m executorch.examples.models.gemma4.webgpu_artifact_manifest \ + stage-runtime --destination-root staged \ + --plain-root "$PLAIN_ROOT" --plain-receipt "$PLAIN_MANIFEST" \ + --mtp-root "$MTP_ROOT" --mtp-receipt "$MTP_MANIFEST" \ + --runtime-source-receipt runtime-source.json \ + --target-prefill-receipt target-prefill.json \ + --plain-wall-javascript "$WALL_BUILD/backends/webgpu/browser_gemma4_plain/webgpu_llama.js" \ + --plain-wall-wasm "$WALL_BUILD/backends/webgpu/browser_gemma4_plain/webgpu_llama.wasm" \ + --plain-profile-javascript "$PROFILE_BUILD/backends/webgpu/browser_gemma4_plain/webgpu_llama.js" \ + --plain-profile-wasm "$PROFILE_BUILD/backends/webgpu/browser_gemma4_plain/webgpu_llama.wasm" \ + --mtp-wall-javascript "$WALL_BUILD/backends/webgpu/browser_gemma4_mtp/gemma4_mtp.js" \ + --mtp-wall-wasm "$WALL_BUILD/backends/webgpu/browser_gemma4_mtp/gemma4_mtp.wasm" \ + --mtp-profile-javascript "$PROFILE_BUILD/backends/webgpu/browser_gemma4_mtp/gemma4_mtp_profile.js" \ + --mtp-profile-wasm "$PROFILE_BUILD/backends/webgpu/browser_gemma4_mtp/gemma4_mtp_profile.wasm" \ + --source-manifest "$SOURCE_MANIFEST" --wgsl-manifest "$WGSL_MANIFEST" \ + --plain-wall-recipe plain-wall-recipe.json \ + --plain-profile-recipe plain-profile-recipe.json \ + --mtp-wall-recipe mtp-wall-recipe.json \ + --mtp-profile-recipe mtp-profile-recipe.json +``` + +This writes a version-3 `gemma4_webgpu_combined_runtime.json` inside a newly +created staging root after validating and copying every referenced byte. +Runtime-source receipt schema version 4 binds both model manifests, the plain +and MTP wall/profile builds, their targets, factories, output stems, profiling +modes, and the source and generated-WGSL manifests. Build-recipe schema version +2 records canonical configure and build arguments for reproduction; +`build_execution: not_attested` explicitly means the receipt does not claim +those commands produced the supplied runtime bytes. +The target-prefill receipt binds the exact runtime-source receipt and reviewed +producer bytes; named-owner checkpoint execution remains required. +Accepted behavior-oracle PTEs, empty runtimes, unbound bytes, symlinks, and +extra staged files fail closed. Plain and MTP views remain pending GPU execution +validation, and the combined view remains pending +cross-view GPU execution validation; source-bound bytes alone convey no +correctness or performance claim. + +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 +graph-owned resources before releasing its owned process WebGPU context. + ### Text ("Write a short paragraph about the history of artificial intelligence") | Model | Size | Load | Prefill | Gen | TTFT | Total | Mem load | Mem peak | diff --git a/examples/models/gemma4/export_speculative.py b/examples/models/gemma4/export_speculative.py index 7a365b7cec5..15dcf32c821 100644 --- a/examples/models/gemma4/export_speculative.py +++ b/examples/models/gemma4/export_speculative.py @@ -10,9 +10,7 @@ from __future__ import annotations import argparse -import json import operator -import shutil import tempfile from pathlib import Path from typing import Any @@ -27,9 +25,8 @@ validate_assistant_checkpoint, ) from executorch.examples.models.gemma4.webgpu_artifact_manifest import ( - create_mtp_manifest, + finalize_mtp_export, validate_export_identity, - validate_mtp_manifest, WEBGPU_EXTERNAL_CONSTANTS_MAX_DATA_BYTES, ) from executorch.examples.models.gemma4.webgpu_partitioner import ( @@ -295,9 +292,7 @@ def export_speculative( # noqa: C901 with tempfile.TemporaryDirectory( prefix=f".{output_path.stem}.", dir=output_path.parent.parent - ) as staging_directory, tempfile.TemporaryDirectory( - prefix=f".{receipt_path.stem}.", dir=receipt_path.parent - ) as receipt_staging_directory: + ) as staging_directory: staging = Path(staging_directory) staged_pte = staging / output_path.name with staged_pte.open("xb") as output: @@ -313,62 +308,23 @@ def export_speculative( # noqa: C901 if staged_pte.stat().st_size == 0: raise ValueError("K=2 export produced an empty PTE") - role_paths: dict[str, Path] = {"pte": staged_pte} - staged_source: Path | None = None - if source_receipt_path is not None: - if source_receipt_path.is_symlink() or not source_receipt_path.is_file(): - raise ValueError("Gemma 4 MTP source receipt must be a regular file") - staged_source = staging / source_receipt_path.name - if staged_source.exists() or staged_source.is_symlink(): - raise ValueError( - "Gemma 4 MTP source receipt basename collides with an artifact" - ) - destination = output_path.parent / staged_source.name - if destination.exists() or destination.is_symlink(): - raise ValueError(f"refusing to overwrite existing artifact: {destination}") - shutil.copyfile(source_receipt_path, staged_source) - role_paths["source"] = staged_source - receipt = create_mtp_manifest(staging, role_paths, staged_tensor_paths) - receipt["evidence"] = { + evidence = { "assistant_checkpoint": assistant_checkpoint_evidence, "k2_abi": k2_abi_evidence, "lowering": lowering_evidence, "qat_selection": qat_selection_evidence, "target_checkpoint": target_checkpoint_evidence, } - validate_mtp_manifest(staging, receipt) - staged_receipt = Path(receipt_staging_directory) / receipt_path.name - staged_receipt.write_text( - json.dumps(receipt, indent=2, sort_keys=True) + "\n", - encoding="utf-8", + return finalize_mtp_export( + staging, + output_path, + receipt_path, + staged_pte, + staged_tensor_paths, + source_receipt_path, + evidence, ) - publications = [ - (path, output_path.parent / path.name) for path in staged_tensor_paths - ] - if staged_source is not None: - publications.append( - (staged_source, output_path.parent / staged_source.name) - ) - publications.append((staged_pte, output_path)) - published: list[Path] = [] - try: - for staged, destination in publications: - staged.replace(destination) - published.append(destination) - validate_mtp_manifest(output_path.parent, receipt) - staged_receipt.replace(receipt_path) - published.append(receipt_path) - final_receipt = json.loads(receipt_path.read_text(encoding="utf-8")) - if not isinstance(final_receipt, dict): - raise ValueError("Gemma 4 MTP receipt must be a JSON object") - validate_mtp_manifest(output_path.parent, final_receipt) - except (OSError, ValueError): - for destination in reversed(published): - destination.unlink(missing_ok=True) - raise - return receipt_path - def main() -> int: parser = argparse.ArgumentParser(description="Export Gemma 4 K=2 for WebGPU") diff --git a/examples/models/gemma4/runner/gemma4_spec_main.cpp b/examples/models/gemma4/runner/gemma4_spec_main.cpp new file mode 100644 index 00000000000..10d7095b1f2 --- /dev/null +++ b/examples/models/gemma4/runner/gemma4_spec_main.cpp @@ -0,0 +1,156 @@ +/* + * 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 + +namespace { + +using ::executorch::examples::gemma4::Gemma4SpecLoadMode; +using ::executorch::examples::gemma4::Gemma4SpecRunner; +using ::executorch::examples::gemma4::Gemma4SpecRunnerConfig; +using ::executorch::examples::gemma4::validate_gemma4_spec_request; +using ::executorch::runtime::Error; + +struct Arguments { + std::string pte; + std::vector ptd; + std::vector prompt_ids; + std::vector stop_tokens; + size_t max_new_tokens = 0; + Gemma4SpecLoadMode load_mode = Gemma4SpecLoadMode::File; +}; + +bool parse_int64(std::string_view value, int64_t* output) { + const char* begin = value.data(); + const char* end = begin + value.size(); + const auto result = std::from_chars(begin, end, *output); + return result.ec == std::errc() && result.ptr == end; +} + +bool parse_size(std::string_view value, size_t* output) { + const char* begin = value.data(); + const char* end = begin + value.size(); + const auto result = std::from_chars(begin, end, *output); + return result.ec == std::errc() && result.ptr == end && *output > 0; +} + +bool parse_token_list(std::string_view value, std::vector* output) { + size_t start = 0; + while (start <= value.size()) { + const size_t end = value.find(',', start); + const std::string_view token = value.substr(start, end - start); + int64_t parsed = -1; + if (token.empty() || !parse_int64(token, &parsed)) { + return false; + } + output->push_back(parsed); + if (end == std::string_view::npos) { + break; + } + start = end + 1; + } + return true; +} + +bool parse_arguments(int argc, char** argv, Arguments* arguments) { + for (int index = 1; index < argc; ++index) { + const std::string_view option(argv[index]); + if (option == "--mmap") { + arguments->load_mode = Gemma4SpecLoadMode::Mmap; + continue; + } + if (index + 1 >= argc) { + return false; + } + const std::string_view value(argv[++index]); + if (option == "--pte") { + arguments->pte = value; + } else if (option == "--ptd") { + arguments->ptd.emplace_back(value); + } else if (option == "--prompt-ids") { + if (!parse_token_list(value, &arguments->prompt_ids)) { + return false; + } + } else if (option == "--stop-token") { + int64_t token = -1; + if (!parse_int64(value, &token)) { + return false; + } + arguments->stop_tokens.push_back(token); + } else if (option == "--max-new-tokens") { + if (!parse_size(value, &arguments->max_new_tokens)) { + return false; + } + } else { + return false; + } + } + return !arguments->pte.empty() && arguments->ptd.size() == 3 && + !arguments->prompt_ids.empty() && arguments->max_new_tokens > 0; +} + +void print_usage(const char* program) { + std::cerr << "Usage: " << program + << " --pte MODEL.pte --ptd A.ptd --ptd B.ptd --ptd C.ptd" + " --prompt-ids ID[,ID...] --max-new-tokens N" + " [--stop-token ID] [--mmap]\n"; +} + +} // namespace + +int main(int argc, char** argv) { + if (argc == 2 && std::string_view(argv[1]) == "--help") { + print_usage(argv[0]); + return 0; + } + Arguments arguments; + if (!parse_arguments(argc, argv, &arguments)) { + print_usage(argv[0]); + return 2; + } + + const Gemma4SpecRunnerConfig config; + if (validate_gemma4_spec_request( + config, + arguments.prompt_ids, + arguments.max_new_tokens, + arguments.stop_tokens) != Error::Ok) { + std::cerr << "Invalid prompt, token budget, stop token, or capacity\n"; + return 3; + } + + Gemma4SpecRunner runner(config); + const Error load_error = + runner.load(arguments.pte, std::move(arguments.ptd), arguments.load_mode); + if (load_error != Error::Ok) { + std::cerr << "Failed to load `k2_round`: " + << static_cast(load_error) << '\n'; + return 4; + } + auto trace = runner.generate( + arguments.prompt_ids, arguments.max_new_tokens, arguments.stop_tokens); + if (!trace.ok()) { + std::cerr << "Generation failed: " << static_cast(trace.error()) + << '\n'; + (void)runner.unload(); + return 5; + } + for (size_t index = 0; index < trace->tokens.size(); ++index) { + std::cout << (index == 0 ? "" : " ") << trace->tokens[index]; + } + std::cout << '\n'; + return runner.unload() == Error::Ok ? 0 : 6; +} diff --git a/examples/models/gemma4/runner/gemma4_spec_runner.cpp b/examples/models/gemma4/runner/gemma4_spec_runner.cpp new file mode 100644 index 00000000000..8d40b0c89f5 --- /dev/null +++ b/examples/models/gemma4/runner/gemma4_spec_runner.cpp @@ -0,0 +1,1003 @@ +/* + * 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 +#include + +namespace executorch::examples::gemma4 { +namespace { + +using ::executorch::aten::ScalarType; +using ::executorch::aten::Tensor; +using ::executorch::backends::webgpu::compare_and_set_default_webgpu_context; +using ::executorch::backends::webgpu::create_webgpu_context; +using ::executorch::backends::webgpu::destroy_webgpu_context; +using ::executorch::backends::webgpu::get_explicit_default_webgpu_context; +using ::executorch::backends::webgpu::load_webgpu_model; +using ::executorch::backends::webgpu::WebGPUContext; +using ::executorch::backends::webgpu::WebGPUModelLoadSpec; +using ::executorch::extension::make_tensor_ptr; +using ::executorch::runtime::Error; +using ::executorch::runtime::EValue; +using ::executorch::runtime::MethodMeta; +using ::executorch::runtime::Result; +using ::executorch::runtime::Tag; +using ::executorch::runtime::TensorInfo; + +bool verify_context(const WebGPUContext* context) { + return context != nullptr && get_explicit_default_webgpu_context() == context; +} + +bool shape_is(const Tensor& tensor, std::initializer_list expected) { + if (tensor.dim() != static_cast(expected.size())) { + return false; + } + size_t index = 0; + for (int32_t dimension : expected) { + if (tensor.size(index++) != dimension) { + return false; + } + } + return true; +} + +bool tensor_shape_is( + const TensorInfo& tensor, + std::initializer_list expected) { + const auto sizes = tensor.sizes(); + if (sizes.size() != expected.size()) { + return false; + } + size_t index = 0; + for (int32_t dimension : expected) { + if (sizes[index++] != dimension) { + return false; + } + } + return true; +} + +bool all_tensor_tags(const MethodMeta& meta) { + for (size_t index = 0; index < meta.num_inputs(); ++index) { + auto tag = meta.input_tag(index); + if (!tag.ok() || tag.get() != Tag::Tensor) { + return false; + } + } + for (size_t index = 0; index < meta.num_outputs(); ++index) { + auto tag = meta.output_tag(index); + if (!tag.ok() || tag.get() != Tag::Tensor) { + return false; + } + } + return true; +} + +bool method_contract_is( + const MethodMeta& meta, + const Gemma4SpecRunnerConfig& config) { + if (std::string_view(meta.name()) != config.method_name || + meta.num_inputs() != 4 || meta.num_outputs() != 5 || + meta.num_backends() != 1 || meta.num_instructions() != 1 || + !all_tensor_tags(meta)) { + return false; + } + auto backend = meta.get_backend_name(0); + if (!backend.ok() || std::string_view(backend.get()) != "VulkanBackend") { + return false; + } + + auto input_ids = meta.input_tensor_meta(0); + auto input_positions = meta.input_tensor_meta(1); + auto is_round = meta.input_tensor_meta(2); + auto donor_length = meta.input_tensor_meta(3); + if (!input_ids.ok() || !input_positions.ok() || !is_round.ok() || + !donor_length.ok() || input_ids->scalar_type() != ScalarType::Long || + input_positions->scalar_type() != ScalarType::Long || + is_round->scalar_type() != ScalarType::Long || + donor_length->scalar_type() != ScalarType::Long || + input_ids->sizes().size() != 2 || input_ids->sizes()[0] != 1 || + input_ids->sizes()[1] != config.max_input_length || + input_positions->sizes().size() != 1 || + input_positions->sizes()[0] != input_ids->sizes()[1] || + !tensor_shape_is(is_round.get(), {1}) || + !tensor_shape_is(donor_length.get(), {1, 1})) { + return false; + } + + auto candidates = meta.output_tensor_meta(0); + auto target_greedy = meta.output_tensor_meta(1); + auto output_matches = meta.output_tensor_meta(2); + auto output_bonus = meta.output_tensor_meta(3); + auto state_probe = meta.output_tensor_meta(4); + return candidates.ok() && target_greedy.ok() && output_matches.ok() && + output_bonus.ok() && state_probe.ok() && + candidates->scalar_type() == ScalarType::Long && + target_greedy->scalar_type() == ScalarType::Long && + output_matches->scalar_type() == ScalarType::Long && + output_bonus->scalar_type() == ScalarType::Long && + state_probe->scalar_type() == ScalarType::Float && + tensor_shape_is(candidates.get(), {1, 2}) && + tensor_shape_is(target_greedy.get(), {1, 3}) && + tensor_shape_is(output_matches.get(), {1}) && + tensor_shape_is(output_bonus.get(), {1, 1}) && + tensor_shape_is(state_probe.get(), {1, 1}); +} + +bool request_fits_capacity( + size_t prompt_depth, + size_t token_budget, + size_t speculative_tail, + int64_t capacity) { + if (capacity <= 0) { + return false; + } + const auto limit = static_cast(capacity); + const auto prompt = static_cast(prompt_depth); + const auto tokens = static_cast(token_budget); + const auto tail = static_cast(speculative_tail); + return prompt <= limit && tokens <= limit - prompt && + tail <= limit - prompt - tokens; +} + +bool position_range_fits_capacity( + int64_t start_position, + size_t count, + int64_t capacity) { + return start_position >= 0 && capacity >= start_position && + static_cast(count) <= + static_cast(capacity - start_position); +} + +#ifdef WGPU_BACKEND_ENABLE_PROFILING +void append_json_string(std::ostringstream& output, const std::string& value) { + output << '"'; + for (const unsigned char character : value) { + switch (character) { + case '"': + output << "\\\""; + break; + case '\\': + output << "\\\\"; + break; + case '\b': + output << "\\b"; + break; + case '\f': + output << "\\f"; + break; + case '\n': + output << "\\n"; + break; + case '\r': + output << "\\r"; + break; + case '\t': + output << "\\t"; + break; + default: + if (character < 0x20) { + constexpr char kHex[] = "0123456789abcdef"; + output << "\\u00" << kHex[character >> 4] << kHex[character & 0x0f]; + } else { + output << static_cast(character); + } + } + } + output << '"'; +} + +std::string serialize_profile_json( + bool timestamp_supported, + bool fresh, + bool valid, + const std::vector<::executorch::backends::webgpu::ShaderDuration>& + durations, + uint64_t execute_generation, + uint64_t context_generation, + uint64_t querypool_generation) { + std::ostringstream output; + output << std::fixed << std::setprecision(9) + << "{\"schemaVersion\":1,\"supported\":" + << (timestamp_supported ? "true" : "false") + << ",\"fresh\":" << (fresh ? "true" : "false") + << ",\"valid\":" << (valid ? "true" : "false") + << ",\"context_generation\":" << context_generation + << ",\"querypool_generation\":" << querypool_generation + << ",\"execute_generation\":" << execute_generation; + if (!timestamp_supported || !fresh || !valid) { + output << ",\"total_kernel_ms\":0,\"pass_span_ms\":0," + "\"interpass_gap_ms\":0,\"perop\":[]}"; + return output.str(); + } + + struct Aggregate { + uint64_t nanoseconds = 0; + uint64_t calls = 0; + }; + std::map per_op; + uint64_t total_nanoseconds = 0; + uint64_t first_start = 0; + uint64_t last_end = 0; + bool have_timestamp = false; + for (const auto& duration : durations) { + auto& aggregate = per_op + [duration.kernel_name.empty() ? std::string("dispatch") + : duration.kernel_name]; + aggregate.nanoseconds += duration.execution_duration_ns; + ++aggregate.calls; + total_nanoseconds += duration.execution_duration_ns; + if (!have_timestamp) { + first_start = duration.start_time_ns; + last_end = duration.end_time_ns; + have_timestamp = true; + } else { + first_start = std::min(first_start, duration.start_time_ns); + last_end = std::max(last_end, duration.end_time_ns); + } + } + const uint64_t span_nanoseconds = + have_timestamp && last_end > first_start ? last_end - first_start : 0; + const uint64_t gap_nanoseconds = span_nanoseconds > total_nanoseconds + ? span_nanoseconds - total_nanoseconds + : 0; + constexpr double kNanosecondsPerMillisecond = 1.0e6; + output << ",\"total_kernel_ms\":" + << total_nanoseconds / kNanosecondsPerMillisecond + << ",\"pass_span_ms\":" + << span_nanoseconds / kNanosecondsPerMillisecond + << ",\"interpass_gap_ms\":" + << gap_nanoseconds / kNanosecondsPerMillisecond << ",\"perop\":["; + bool first = true; + for (const auto& [name, aggregate] : per_op) { + if (!first) { + output << ','; + } + first = false; + output << "{\"op\":"; + append_json_string(output, name); + output << ",\"ms\":" << aggregate.nanoseconds / kNanosecondsPerMillisecond + << ",\"calls\":" << aggregate.calls << '}'; + } + output << "]}"; + return output.str(); +} +#endif + +#ifndef WGPU_BACKEND_ENABLE_PROFILING +std::string unsupported_profile_json( + uint64_t execute_generation, + uint64_t context_generation) { + return "{\"schemaVersion\":1,\"supported\":false," + "\"fresh\":false,\"valid\":false," + "\"context_generation\":" + + std::to_string(context_generation) + + ",\"querypool_generation\":0," + "\"execute_generation\":" + + std::to_string(execute_generation) + + ",\"total_kernel_ms\":0,\"pass_span_ms\":0," + "\"interpass_gap_ms\":0,\"perop\":[]}"; +} +#endif + +} // namespace + +Error validate_gemma4_spec_request( + const Gemma4SpecRunnerConfig& config, + const std::vector& prompt_ids, + size_t token_budget, + const std::vector& stop_tokens) { + constexpr size_t kSpeculativeTail = 2; + if (config.vocab_size <= 0 || config.max_input_length <= 0 || + config.max_input_length > std::numeric_limits::max() || + config.target_capacity <= 0 || config.donor_capacity <= 0 || + prompt_ids.empty() || token_budget == 0) { + return Error::InvalidArgument; + } + const auto valid_token = [&config](int64_t token) { + return token >= 0 && token < config.vocab_size; + }; + if (!std::all_of(prompt_ids.begin(), prompt_ids.end(), valid_token) || + !std::all_of(stop_tokens.begin(), stop_tokens.end(), valid_token)) { + return Error::InvalidArgument; + } + + const size_t speculative_tail = token_budget > 1 ? kSpeculativeTail : 0; + if ((token_budget > 1 && prompt_ids.size() < 2) || + !request_fits_capacity( + prompt_ids.size(), + token_budget, + speculative_tail, + config.target_capacity) || + !request_fits_capacity( + prompt_ids.size(), + token_budget, + speculative_tail, + config.donor_capacity)) { + return Error::InvalidArgument; + } + return Error::Ok; +} + +class Gemma4SpecRunner::Impl final { + public: + explicit Impl(Gemma4SpecRunnerConfig runner_config) + : config(std::move(runner_config)) {} + + bool valid_token(int64_t token) const { + return token >= 0 && token < config.vocab_size; + } + + void clear_controller_state() { + buffered.clear(); + incremental_prefill.clear(); + execute_count = 0; + accepted_drafts = 0; + next_position = -1; + last_emitted = -1; + } + + void arm_profile() { + consumed_profile_generation = profile_generation; + } + + void clear_profile_binding() { +#ifdef WGPU_BACKEND_ENABLE_PROFILING + has_profile_binding = false; + bound_profile_generation = 0; + bound_querypool_generation = 0; +#endif + } + + bool enable_profile_environment() { + if (profile_environment_owned) { + return true; + } + const char* current = std::getenv("WEBGPU_TIMESTAMP_QUERY"); + previous_profile_environment = + current != nullptr ? std::optional(current) : std::nullopt; + if (setenv("WEBGPU_TIMESTAMP_QUERY", "1", 1) != 0) { + previous_profile_environment.reset(); + return false; + } + profile_environment_owned = true; + return true; + } + + void restore_profile_environment() { + if (!profile_environment_owned) { + return; + } + if (previous_profile_environment.has_value()) { + (void)setenv( + "WEBGPU_TIMESTAMP_QUERY", previous_profile_environment->c_str(), 1); + } else { + (void)unsetenv("WEBGPU_TIMESTAMP_QUERY"); + } + previous_profile_environment.reset(); + profile_environment_owned = false; + } + + Gemma4SpecRunnerConfig config; + std::unique_ptr context; + std::unique_ptr module; + std::deque buffered; + std::vector incremental_prefill; + size_t execute_count = 0; + size_t accepted_drafts = 0; + int64_t next_position = -1; + int64_t last_emitted = -1; + bool method_fresh = false; + bool method_healthy = false; + bool profile_enabled = false; + bool profile_environment_owned = false; + std::optional previous_profile_environment; + uint64_t context_generation = 0; + uint64_t profile_generation = 0; + uint64_t consumed_profile_generation = 0; +#ifdef WGPU_BACKEND_ENABLE_PROFILING + bool has_profile_binding = false; + uint64_t bound_profile_generation = 0; + uint64_t bound_querypool_generation = 0; +#endif +}; + +Gemma4SpecRunner::Gemma4SpecRunner(Gemma4SpecRunnerConfig config) + : impl_(std::make_unique(std::move(config))) {} + +Gemma4SpecRunner::~Gemma4SpecRunner() { + (void)unload(); +} + +Error Gemma4SpecRunner::load( + const std::string& pte_path, + std::vector ptd_paths, + Gemma4SpecLoadMode load_mode) { + if (pte_path.empty() || impl_->config.vocab_size <= 0 || + impl_->config.max_input_length <= 0 || + impl_->config.max_input_length > std::numeric_limits::max() || + impl_->config.target_capacity <= 0 || impl_->config.donor_capacity <= 0 || + impl_->config.method_name != "k2_round" || ptd_paths.size() != 3) { + return Error::InvalidArgument; + } + + const bool acquired_context = impl_->context == nullptr; + if (acquired_context) { + try { + impl_->context = std::make_unique(create_webgpu_context()); + } catch (const std::exception&) { + impl_->context.reset(); + return Error::Internal; + } + if (!compare_and_set_default_webgpu_context( + nullptr, impl_->context.get())) { + destroy_webgpu_context(*impl_->context); + impl_->context.reset(); + impl_->method_fresh = false; + impl_->method_healthy = false; + return Error::InvalidState; + } + } else if (!verify_context(impl_->context.get())) { + impl_->method_fresh = false; + impl_->method_healthy = false; + return Error::InvalidState; + } + const auto release_acquired_context = [&]() { + if (!acquired_context) { + return true; + } + const bool released = + compare_and_set_default_webgpu_context(impl_->context.get(), nullptr); + destroy_webgpu_context(*impl_->context); + impl_->context.reset(); + return released; + }; + + WebGPUModelLoadSpec spec; + spec.pte_path = pte_path; + spec.ptd_paths = std::move(ptd_paths); + spec.required_methods = {impl_->config.method_name}; + spec.load_mode = load_mode == Gemma4SpecLoadMode::File + ? extension::Module::LoadMode::File + : extension::Module::LoadMode::Mmap; + auto loaded = load_webgpu_model(std::move(spec)); + if (!loaded.ok()) { + return release_acquired_context() ? loaded.error() : Error::InvalidState; + } + + auto next_module = std::move(loaded.get()); + auto methods = next_module->method_names(); + auto meta = next_module->method_meta(impl_->config.method_name); + if (!methods.ok() || methods->size() != 1 || + methods->count(impl_->config.method_name) != 1 || !meta.ok() || + !method_contract_is(meta.get(), impl_->config)) { + if (next_module->is_method_loaded(impl_->config.method_name)) { + (void)next_module->unload_method(impl_->config.method_name); + } + next_module.reset(); + return release_acquired_context() ? Error::InvalidProgram + : Error::InvalidState; + } + if (!verify_context(impl_->context.get())) { + (void)next_module->unload_method(impl_->config.method_name); + next_module.reset(); + (void)release_acquired_context(); + return Error::InvalidState; + } + + if (impl_->module != nullptr && + impl_->module->is_method_loaded(impl_->config.method_name) && + !impl_->module->unload_method(impl_->config.method_name)) { + (void)next_module->unload_method(impl_->config.method_name); + impl_->method_fresh = false; + impl_->method_healthy = false; + return Error::Internal; + } + impl_->module = std::move(next_module); + impl_->clear_controller_state(); + impl_->method_fresh = true; + impl_->method_healthy = true; + impl_->profile_enabled = false; + if (acquired_context) { + ++impl_->context_generation; + } + impl_->profile_generation = 0; + impl_->arm_profile(); + impl_->clear_profile_binding(); + impl_->restore_profile_environment(); + return Error::Ok; +} + +Error Gemma4SpecRunner::reset() { + if (impl_->module == nullptr || + !impl_->module->is_method_loaded(impl_->config.method_name) || + impl_->context == nullptr || !verify_context(impl_->context.get())) { + impl_->method_healthy = false; + impl_->method_fresh = false; + impl_->clear_controller_state(); + return Error::InvalidState; + } + impl_->arm_profile(); + impl_->clear_profile_binding(); + impl_->clear_controller_state(); + if (!impl_->module->unload_method(impl_->config.method_name)) { + impl_->method_healthy = false; + impl_->method_fresh = false; + return Error::Internal; + } + const Error error = impl_->module->load_method(impl_->config.method_name); + impl_->method_healthy = + error == Error::Ok && verify_context(impl_->context.get()); + impl_->method_fresh = impl_->method_healthy; + return error != Error::Ok + ? error + : (impl_->method_healthy ? Error::Ok : Error::InvalidState); +} + +Error Gemma4SpecRunner::unload() { + if (impl_->context != nullptr && !verify_context(impl_->context.get())) { + return Error::InvalidState; + } + impl_->profile_enabled = false; + impl_->restore_profile_environment(); + bool method_unloaded = true; + if (impl_->module != nullptr) { + if (impl_->module->is_method_loaded(impl_->config.method_name)) { + method_unloaded = impl_->module->unload_method(impl_->config.method_name); + } + impl_->module.reset(); + } + impl_->clear_profile_binding(); + impl_->method_fresh = false; + impl_->method_healthy = false; + impl_->clear_controller_state(); + if (impl_->context != nullptr) { + if (!compare_and_set_default_webgpu_context( + impl_->context.get(), nullptr)) { + return Error::InvalidState; + } + destroy_webgpu_context(*impl_->context); + impl_->context.reset(); + } + return method_unloaded ? Error::Ok : Error::Internal; +} + +bool Gemma4SpecRunner::is_loaded() const { + return impl_->module != nullptr && impl_->method_healthy && + impl_->context != nullptr && verify_context(impl_->context.get()); +} + +Result Gemma4SpecRunner::execute( + const std::vector& input_ids, + const std::vector& input_positions, + bool is_round, + int64_t donor_length) { + if (!is_loaded() || impl_->context == nullptr || + !verify_context(impl_->context.get())) { + impl_->method_healthy = false; + return Error::InvalidState; + } + if (input_ids.empty() || input_ids.size() != input_positions.size() || + donor_length < 0 || donor_length > impl_->config.donor_capacity || + !position_range_fits_capacity( + input_positions.front(), + input_positions.size(), + impl_->config.target_capacity) || + !position_range_fits_capacity( + input_positions.front(), + input_positions.size(), + impl_->config.donor_capacity)) { + return Error::InvalidArgument; + } + if ((!is_round && + input_ids.size() > + static_cast(impl_->config.max_input_length)) || + (is_round && input_ids.size() != 3)) { + return Error::InvalidArgument; + } + for (size_t index = 0; index < input_ids.size(); ++index) { + if (!impl_->valid_token(input_ids[index]) || input_positions[index] < 0 || + (index > 0 && + input_positions[index] != input_positions[index - 1] + 1)) { + return Error::InvalidArgument; + } + } + const int64_t start_position = input_positions.front(); + if ((is_round && (donor_length != start_position || donor_length < 2)) || + (!is_round && start_position == 0 && donor_length != 2) || + (!is_round && start_position > 0 && donor_length != start_position)) { + return Error::InvalidArgument; + } + + auto ids = + make_tensor_ptr({1, static_cast(input_ids.size())}, input_ids); + auto positions = make_tensor_ptr( + {static_cast(input_positions.size())}, input_positions); + auto round = make_tensor_ptr({1}, std::vector{is_round ? 1 : 0}); + auto donor = make_tensor_ptr({1, 1}, std::vector{donor_length}); + + impl_->clear_profile_binding(); + impl_->method_fresh = false; + auto execution = impl_->module->execute( + impl_->config.method_name, + {EValue(ids), EValue(positions), EValue(round), EValue(donor)}); + if (!verify_context(impl_->context.get())) { + impl_->method_healthy = false; + return Error::InvalidState; + } + if (!execution.ok()) { + impl_->method_healthy = false; + return execution.error(); + } + if (execution->size() != 5) { + impl_->method_healthy = false; + return Error::InvalidProgram; + } + for (const EValue& value : *execution) { + if (!value.isTensor()) { + impl_->method_healthy = false; + return Error::InvalidType; + } + } + + const Tensor& candidates = execution->at(0).toTensor(); + const Tensor& target = execution->at(1).toTensor(); + const Tensor& matches = execution->at(2).toTensor(); + const Tensor& bonus = execution->at(3).toTensor(); + const Tensor& probe = execution->at(4).toTensor(); + if (candidates.scalar_type() != ScalarType::Long || + target.scalar_type() != ScalarType::Long || + matches.scalar_type() != ScalarType::Long || + bonus.scalar_type() != ScalarType::Long || + probe.scalar_type() != ScalarType::Float || + !shape_is(candidates, {1, 2}) || !shape_is(target, {1, 3}) || + !shape_is(matches, {1}) || !shape_is(bonus, {1, 1}) || + !shape_is(probe, {1, 1})) { + impl_->method_healthy = false; + return Error::InvalidProgram; + } + + Gemma4K2Output output; + const int64_t* candidate_data = candidates.const_data_ptr(); + const int64_t* target_data = target.const_data_ptr(); + output.candidates = {candidate_data[0], candidate_data[1]}; + output.target_greedy = {target_data[0], target_data[1], target_data[2]}; + output.match_count = matches.const_data_ptr()[0]; + output.bonus = bonus.const_data_ptr()[0]; + output.state_probe = probe.const_data_ptr()[0]; + + for (int64_t token : output.candidates) { + if (!impl_->valid_token(token)) { + impl_->method_healthy = false; + return Error::InvalidProgram; + } + } + for (int64_t token : output.target_greedy) { + if (!impl_->valid_token(token)) { + impl_->method_healthy = false; + return Error::InvalidProgram; + } + } + + if (is_round) { + const auto decision = reconcile_gemma4_k2( + output, start_position, 3, {}, impl_->config.vocab_size); + if (!decision.valid) { + impl_->method_healthy = false; + return Error::InvalidProgram; + } + } else if ( + output.match_count != 0 || + output.target_greedy[0] != output.target_greedy[1] || + output.target_greedy[1] != output.target_greedy[2] || + output.bonus != output.target_greedy[0] || + !impl_->valid_token(output.bonus) || !std::isfinite(output.state_probe)) { + impl_->method_healthy = false; + return Error::InvalidProgram; + } + ++impl_->execute_count; + ++impl_->profile_generation; +#ifdef WGPU_BACKEND_ENABLE_PROFILING + if (impl_->profile_enabled && impl_->context->timestamp_supported && + impl_->context->querypool != nullptr) { + impl_->has_profile_binding = true; + impl_->bound_profile_generation = impl_->profile_generation; + impl_->bound_querypool_generation = + impl_->context->querypool->result_generation(); + } +#endif + return output; +} + +Result Gemma4SpecRunner::prefill( + const std::vector& input_ids, + int64_t start_position) { + if (input_ids.empty() || + input_ids.size() > static_cast(impl_->config.max_input_length) || + !position_range_fits_capacity( + start_position, input_ids.size(), impl_->config.target_capacity) || + !position_range_fits_capacity( + start_position, input_ids.size(), impl_->config.donor_capacity) || + !std::all_of(input_ids.begin(), input_ids.end(), [this](int64_t token) { + return impl_->valid_token(token); + })) { + return Error::InvalidArgument; + } + if (start_position == 0) { + impl_->buffered.clear(); + impl_->incremental_prefill.clear(); + if (!impl_->method_fresh) { + const Error error = reset(); + if (error != Error::Ok) { + return error; + } + } + } else if (start_position != impl_->next_position) { + return Error::InvalidArgument; + } + std::vector positions; + positions.reserve(input_ids.size()); + for (size_t index = 0; index < input_ids.size(); ++index) { + positions.push_back(start_position + static_cast(index)); + } + auto output = execute( + input_ids, positions, false, start_position == 0 ? 2 : start_position); + if (!output.ok()) { + return output.error(); + } + impl_->next_position = + start_position + static_cast(input_ids.size()); + impl_->last_emitted = output->bonus; + return output->bonus; +} + +Error Gemma4SpecRunner::prefill_step(int64_t token, int64_t position) { + const int64_t max_incremental_position = std::min( + {impl_->config.max_input_length, + impl_->config.target_capacity, + impl_->config.donor_capacity}); + if (!is_loaded() || !impl_->valid_token(token) || position < 0 || + max_incremental_position <= 1 || + position >= max_incremental_position - 1) { + return Error::InvalidArgument; + } + if (position == 0) { + impl_->incremental_prefill.clear(); + } + if (static_cast(position) != impl_->incremental_prefill.size()) { + impl_->incremental_prefill.clear(); + return Error::InvalidArgument; + } + impl_->incremental_prefill.push_back(token); + return Error::Ok; +} + +Result Gemma4SpecRunner::step( + int64_t seed_token, + int64_t seed_position) { + const int64_t capacity = + std::min(impl_->config.target_capacity, impl_->config.donor_capacity); + if (!is_loaded() || !impl_->valid_token(seed_token) || seed_position < 1 || + seed_position >= capacity) { + return Error::InvalidArgument; + } + if (!impl_->incremental_prefill.empty()) { + if (static_cast(seed_position) != + impl_->incremental_prefill.size() || + seed_position >= impl_->config.max_input_length) { + impl_->incremental_prefill.clear(); + return Error::InvalidArgument; + } + impl_->incremental_prefill.push_back(seed_token); + auto prompt = std::move(impl_->incremental_prefill); + impl_->incremental_prefill.clear(); + return prefill(prompt, 0); + } + if (seed_position != impl_->next_position || + seed_token != impl_->last_emitted) { + return Error::InvalidArgument; + } + if (!impl_->buffered.empty()) { + const int64_t token = impl_->buffered.front(); + impl_->buffered.pop_front(); + impl_->last_emitted = token; + ++impl_->next_position; + return token; + } + + if (!position_range_fits_capacity( + seed_position, 3, impl_->config.target_capacity) || + !position_range_fits_capacity( + seed_position, 3, impl_->config.donor_capacity)) { + return Error::InvalidArgument; + } + + auto output = execute( + {seed_token, 0, 0}, + {seed_position, seed_position + 1, seed_position + 2}, + true, + seed_position); + if (!output.ok()) { + return output.error(); + } + const auto decision = reconcile_gemma4_k2( + output.get(), seed_position, 3, {}, impl_->config.vocab_size); + if (!decision.valid || decision.committed.empty()) { + impl_->method_healthy = false; + return Error::InvalidProgram; + } + impl_->accepted_drafts += decision.accepted_drafts; + for (size_t index = 1; index < decision.committed.size(); ++index) { + impl_->buffered.push_back(decision.committed[index]); + } + impl_->last_emitted = decision.committed.front(); + ++impl_->next_position; + return decision.committed.front(); +} + +Result Gemma4SpecRunner::generate( + const std::vector& prompt_ids, + size_t token_budget, + const std::vector& stop_tokens) { + const Error request_error = validate_gemma4_spec_request( + impl_->config, prompt_ids, token_budget, stop_tokens); + if (request_error != Error::Ok) { + return request_error; + } + + int64_t start_position = 0; + int64_t prefill_token = -1; + while (start_position < static_cast(prompt_ids.size())) { + const int64_t count = std::min( + impl_->config.max_input_length, + static_cast(prompt_ids.size()) - start_position); + std::vector chunk( + prompt_ids.begin() + start_position, + prompt_ids.begin() + start_position + count); + auto result = prefill(chunk, start_position); + if (!result.ok()) { + return result.error(); + } + prefill_token = result.get(); + start_position += count; + } + + Gemma4SpecTrace trace; + trace.prefill_token = prefill_token; + if (std::find(stop_tokens.begin(), stop_tokens.end(), prefill_token) != + stop_tokens.end()) { + trace.stop_token = prefill_token; + } else { + trace.tokens.push_back(prefill_token); + } + + int64_t seed = prefill_token; + while (!trace.stop_token.has_value() && trace.tokens.size() < token_budget) { + const int64_t round_start = start_position; + auto output = execute( + {seed, 0, 0}, + {round_start, round_start + 1, round_start + 2}, + true, + round_start); + if (!output.ok()) { + return output.error(); + } + const auto decision = reconcile_gemma4_k2( + output.get(), + round_start, + token_budget - trace.tokens.size(), + stop_tokens, + impl_->config.vocab_size); + if (!decision.valid) { + impl_->method_healthy = false; + return Error::InvalidProgram; + } + trace.tokens.insert( + trace.tokens.end(), + decision.committed.begin(), + decision.committed.end()); + trace.discarded_tokens += decision.discarded.size(); + trace.accepted_drafts += decision.accepted_drafts; + trace.rounds.push_back(decision); + if (decision.stopped) { + trace.stop_token = decision.stop_token; + } + start_position = decision.next_position; + seed = decision.next_seed; + } + impl_->accepted_drafts += trace.accepted_drafts; + impl_->next_position = start_position; + impl_->last_emitted = seed; + trace.execute_count = impl_->execute_count; + return trace; +} + +void Gemma4SpecRunner::set_profiling_enabled(bool enabled) { + impl_->arm_profile(); + if (impl_->profile_enabled != enabled) { + impl_->clear_profile_binding(); + } +#ifdef WGPU_BACKEND_ENABLE_PROFILING + if (enabled) { + impl_->profile_enabled = impl_->enable_profile_environment(); + } else { + impl_->profile_enabled = false; + impl_->restore_profile_environment(); + } +#else + (void)enabled; + impl_->profile_enabled = false; + impl_->restore_profile_environment(); +#endif +} + +std::string Gemma4SpecRunner::profile_json() { +#ifdef WGPU_BACKEND_ENABLE_PROFILING + const bool supported = + impl_->context != nullptr && impl_->context->timestamp_supported; + const auto* querypool = + impl_->context != nullptr ? impl_->context->querypool.get() : nullptr; + const uint64_t querypool_generation = + querypool != nullptr ? querypool->result_generation() : 0; + const bool binding_current = impl_->profile_enabled && + impl_->has_profile_binding && + impl_->bound_profile_generation == impl_->profile_generation && + impl_->bound_querypool_generation == querypool_generation; + const bool fresh = binding_current && + impl_->consumed_profile_generation != impl_->profile_generation; + const bool valid = + binding_current && querypool != nullptr && querypool->results_valid(); + if (fresh) { + impl_->consumed_profile_generation = impl_->profile_generation; + } + const std::vector<::executorch::backends::webgpu::ShaderDuration> + empty_durations; + const auto& durations = + querypool != nullptr ? querypool->results() : empty_durations; + return serialize_profile_json( + supported, + fresh, + valid, + durations, + impl_->profile_generation, + impl_->context_generation, + querypool_generation); +#else + return unsupported_profile_json( + impl_->profile_generation, impl_->context_generation); +#endif +} + +size_t Gemma4SpecRunner::execute_count() const { + return impl_->execute_count; +} + +size_t Gemma4SpecRunner::accepted_drafts() const { + return impl_->accepted_drafts; +} + +size_t Gemma4SpecRunner::buffered_tokens() const { + return impl_->buffered.size(); +} + +} // namespace executorch::examples::gemma4 diff --git a/examples/models/gemma4/runner/gemma4_spec_runner.h b/examples/models/gemma4/runner/gemma4_spec_runner.h new file mode 100644 index 00000000000..36648d6cc13 --- /dev/null +++ b/examples/models/gemma4/runner/gemma4_spec_runner.h @@ -0,0 +1,185 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace executorch::examples::gemma4 { + +struct Gemma4K2Output { + std::array candidates{}; + std::array target_greedy{}; + int64_t match_count = -1; + int64_t bonus = -1; + float state_probe = 0.0f; +}; + +struct Gemma4K2Decision { + bool valid = false; + bool stopped = false; + int64_t stop_token = -1; + int64_t next_position = -1; + int64_t next_seed = -1; + size_t accepted_drafts = 0; + std::vector selected; + std::vector committed; + std::vector discarded; +}; + +inline Gemma4K2Decision reconcile_gemma4_k2( + const Gemma4K2Output& output, + int64_t start_position, + size_t token_budget, + const std::vector& stop_tokens, + int64_t vocab_size = 262144) { + Gemma4K2Decision decision; + if (start_position < 2 || token_budget == 0 || vocab_size <= 0 || + output.match_count < 0 || output.match_count > 2 || + !std::isfinite(output.state_probe)) { + return decision; + } + const auto valid_token = [vocab_size](int64_t token) { + return token >= 0 && token < vocab_size; + }; + for (int64_t token : output.candidates) { + if (!valid_token(token)) { + return decision; + } + } + for (int64_t token : output.target_greedy) { + if (!valid_token(token)) { + return decision; + } + } + int64_t expected_matches = 0; + if (output.candidates[0] == output.target_greedy[0]) { + expected_matches = output.candidates[1] == output.target_greedy[1] ? 2 : 1; + } + if (output.match_count != expected_matches || !valid_token(output.bonus) || + output.bonus != output.target_greedy[output.match_count]) { + return decision; + } + + decision.accepted_drafts = static_cast(output.match_count); + decision.next_position = start_position + output.match_count + 1; + decision.next_seed = output.bonus; + for (int64_t index = 0; index < output.match_count; ++index) { + decision.selected.push_back(output.candidates[index]); + } + decision.selected.push_back(output.bonus); + + for (size_t index = 0; index < decision.selected.size(); ++index) { + const int64_t token = decision.selected[index]; + if (std::find(stop_tokens.begin(), stop_tokens.end(), token) != + stop_tokens.end()) { + decision.stopped = true; + decision.stop_token = token; + decision.discarded.insert( + decision.discarded.end(), + decision.selected.begin() + index + 1, + decision.selected.end()); + break; + } + if (decision.committed.size() == token_budget) { + decision.discarded.insert( + decision.discarded.end(), + decision.selected.begin() + index, + decision.selected.end()); + break; + } + decision.committed.push_back(token); + } + decision.valid = true; + return decision; +} + +struct Gemma4SpecRunnerConfig { + 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"; +}; + +runtime::Error validate_gemma4_spec_request( + const Gemma4SpecRunnerConfig& config, + const std::vector& prompt_ids, + size_t token_budget, + const std::vector& stop_tokens); + +enum class Gemma4SpecLoadMode { + File, + Mmap, +}; + +struct Gemma4SpecTrace { + int64_t prefill_token = -1; + std::optional stop_token; + size_t execute_count = 0; + size_t accepted_drafts = 0; + size_t discarded_tokens = 0; + std::vector tokens; + std::vector rounds; +}; + +class Gemma4SpecRunner final { + public: + explicit Gemma4SpecRunner(Gemma4SpecRunnerConfig config = {}); + ~Gemma4SpecRunner(); + + Gemma4SpecRunner(const Gemma4SpecRunner&) = delete; + Gemma4SpecRunner& operator=(const Gemma4SpecRunner&) = delete; + + runtime::Error load( + const std::string& pte_path, + std::vector ptd_paths, + Gemma4SpecLoadMode load_mode = Gemma4SpecLoadMode::File); + runtime::Error reset(); + runtime::Error unload(); + + bool is_loaded() const; + runtime::Result execute( + const std::vector& input_ids, + const std::vector& input_positions, + bool is_round, + int64_t donor_length); + runtime::Result prefill( + const std::vector& input_ids, + int64_t start_position); + runtime::Error prefill_step(int64_t token, int64_t position); + runtime::Result step(int64_t seed_token, int64_t seed_position); + runtime::Result generate( + const std::vector& prompt_ids, + size_t token_budget, + const std::vector& stop_tokens); + + void set_profiling_enabled(bool enabled); + std::string profile_json(); + + size_t execute_count() const; + size_t accepted_drafts() const; + size_t buffered_tokens() const; + + private: + class Impl; + std::unique_ptr impl_; +}; + +} // namespace executorch::examples::gemma4 diff --git a/examples/models/gemma4/runner/gemma4_spec_wasm.cpp b/examples/models/gemma4/runner/gemma4_spec_wasm.cpp new file mode 100644 index 00000000000..4aeec06ecf6 --- /dev/null +++ b/examples/models/gemma4/runner/gemma4_spec_wasm.cpp @@ -0,0 +1,194 @@ +/* + * 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 + +#if defined(__EMSCRIPTEN__) +#include +#define ET_WASM_EXPORT EMSCRIPTEN_KEEPALIVE +#else +#define ET_WASM_EXPORT +#endif + +namespace { + +using ::executorch::backends::webgpu::webgpu_backend_execution_attestation_json; +using ::executorch::examples::gemma4::Gemma4SpecRunner; +using ::executorch::runtime::Error; + +constexpr const char* kMethodName = "k2_round"; +constexpr size_t kExpectedTensorDataPaths = 3; + +std::unique_ptr runner; +std::string execution_attestation_json; +std::string profile_json; + +std::vector split_paths(const char* paths) { + std::vector result; + if (paths == nullptr) { + return result; + } + const std::string value(paths); + size_t start = 0; + while (start <= value.size()) { + const size_t end = value.find('\n', start); + const std::string path = value.substr(start, end - start); + if (!path.empty()) { + result.push_back(path); + } + if (end == std::string::npos) { + break; + } + start = end + 1; + } + return result; +} + +} // namespace + +extern "C" { + +ET_WASM_EXPORT int et_init() { + if (runner == nullptr) { + runner = std::make_unique(); + } + return 1; +} + +ET_WASM_EXPORT int et_load( + const char* pte_path, + const char* tensor_data_paths, + const char* method) { + if (runner == nullptr || pte_path == nullptr || method == nullptr || + std::string(method) != kMethodName) { + return 0; + } + auto paths = split_paths(tensor_data_paths); + if (paths.size() != kExpectedTensorDataPaths) { + return 0; + } + return runner->load(pte_path, std::move(paths)) == Error::Ok ? 1 : 0; +} + +ET_WASM_EXPORT int et_unload() { + if (runner == nullptr) { + return 1; + } + const Error error = runner->unload(); + runner.reset(); + return error == Error::Ok ? 1 : 0; +} + +ET_WASM_EXPORT int et_reset() { + return runner != nullptr && runner->reset() == Error::Ok ? 1 : 0; +} + +ET_WASM_EXPORT int +et_prefill_batch(const int32_t* ids, int32_t count, int32_t start_position) { + const int64_t actual_count = count < 0 ? -static_cast(count) : count; + if (runner == nullptr || ids == nullptr || actual_count <= 0 || + actual_count > INT32_MAX || start_position < 0) { + return -1; + } + std::vector input_ids(ids, ids + actual_count); + auto result = runner->prefill(input_ids, start_position); + return result.ok() ? static_cast(result.get()) : -2; +} + +ET_WASM_EXPORT void et_prefill_step(int32_t token, int32_t position) { + if (runner != nullptr) { + (void)runner->prefill_step(token, position); + } +} + +ET_WASM_EXPORT int et_step(int32_t token, int32_t position) { + if (runner == nullptr) { + return -1; + } + auto result = runner->step(token, position); + return result.ok() ? static_cast(result.get()) : -2; +} + +ET_WASM_EXPORT int et_mtp_execute_count() { + return runner == nullptr ? 0 : static_cast(runner->execute_count()); +} + +ET_WASM_EXPORT int et_mtp_accepted_drafts() { + return runner == nullptr ? 0 : static_cast(runner->accepted_drafts()); +} + +ET_WASM_EXPORT int et_mtp_buffered_tokens() { + return runner == nullptr ? 0 : static_cast(runner->buffered_tokens()); +} + +ET_WASM_EXPORT int et_mtp_execute( + const int32_t* ids, + int32_t count, + int32_t start_position, + int32_t is_round, + int32_t donor_length, + int32_t* output) { + if (runner == nullptr || ids == nullptr || output == nullptr || count <= 0 || + start_position < 0 || (is_round != 0 && is_round != 1)) { + return 0; + } + std::vector input_ids(ids, ids + count); + std::vector positions; + positions.reserve(count); + for (int32_t index = 0; index < count; ++index) { + positions.push_back(static_cast(start_position) + index); + } + auto result = + runner->execute(input_ids, positions, is_round != 0, donor_length); + if (!result.ok()) { + return 0; + } + output[0] = static_cast(result->candidates[0]); + output[1] = static_cast(result->candidates[1]); + output[2] = static_cast(result->target_greedy[0]); + output[3] = static_cast(result->target_greedy[1]); + output[4] = static_cast(result->target_greedy[2]); + output[5] = static_cast(result->match_count); + output[6] = static_cast(result->bonus); + return 1; +} + +ET_WASM_EXPORT const char* et_mtp_execution_attestation() { + execution_attestation_json = webgpu_backend_execution_attestation_json(); + return execution_attestation_json.c_str(); +} + +ET_WASM_EXPORT void et_profile_enable(int enabled) { + if (runner != nullptr) { + runner->set_profiling_enabled(enabled != 0); + } +} + +ET_WASM_EXPORT const char* et_profile() { + if (runner == nullptr) { + static const std::string unsupported = + "{\"schemaVersion\":1,\"supported\":false," + "\"fresh\":false,\"valid\":false,\"context_generation\":0," + "\"querypool_generation\":0,\"execute_generation\":0," + "\"total_kernel_ms\":0," + "\"pass_span_ms\":0,\"interpass_gap_ms\":0,\"perop\":[]}"; + return unsupported.c_str(); + } + profile_json = runner->profile_json(); + return profile_json.c_str(); +} + +} // extern "C" diff --git a/examples/models/gemma4/targets.bzl b/examples/models/gemma4/targets.bzl index 23297e95e87..d0079270aae 100644 --- a/examples/models/gemma4/targets.bzl +++ b/examples/models/gemma4/targets.bzl @@ -38,6 +38,7 @@ def define_webgpu_python_targets(): "//executorch/backends/vulkan/partitioner:vulkan_partitioner", "//executorch/backends/vulkan/patterns:vulkan_patterns", "//executorch/backends/webgpu/scripts:webgpu_artifact_manifest", + "//executorch/examples/models/gemma4:target_prefill_contract", "//executorch/exir:lib", ], ) @@ -87,6 +88,42 @@ def define_common_targets(): ], ) + runtime.cxx_library( + name = "gemma4_spec_runner", + srcs = ["runner/gemma4_spec_runner.cpp"], + exported_headers = ["runner/gemma4_spec_runner.h"], + compiler_flags = ["-fexceptions"], + visibility = ["PUBLIC"], + exported_deps = [ + "//executorch/runtime/core:core", + ], + deps = [ + "//executorch/backends/webgpu:webgpu_backend", + "//executorch/backends/webgpu:webgpu_model_loader", + "//executorch/extension/module:module", + "//executorch/extension/tensor:tensor", + ], + ) + + runtime.cxx_library( + name = "gemma4_spec_wasm_adapter", + srcs = ["runner/gemma4_spec_wasm.cpp"], + compiler_flags = ["-fexceptions"], + visibility = ["PUBLIC"], + deps = [ + ":gemma4_spec_runner", + "//executorch/backends/webgpu:webgpu_backend", + ], + ) + + runtime.cxx_binary( + name = "gemma4_spec_runner_cli", + srcs = ["runner/gemma4_spec_main.cpp"], + compiler_flags = ["-fexceptions"], + visibility = ["PUBLIC"], + deps = [":gemma4_spec_runner"], + ) + runtime.cxx_binary( name = "main", srcs = ["e2e_runner.cpp"], diff --git a/examples/models/gemma4/webgpu_artifact_manifest.py b/examples/models/gemma4/webgpu_artifact_manifest.py index fbd7ff5587b..7bc73fab6cf 100644 --- a/examples/models/gemma4/webgpu_artifact_manifest.py +++ b/examples/models/gemma4/webgpu_artifact_manifest.py @@ -10,18 +10,33 @@ import argparse import copy +import dataclasses +import errno import hashlib import importlib.util import json +import math +import os +import shutil +import signal +import stat import subprocess +import tempfile +import threading from pathlib import Path +from types import FrameType from typing import Any, Mapping, Sequence from executorch.backends.webgpu.scripts.webgpu_artifact_manifest import ( create_manifest, validate_manifest, ) +from executorch.examples.models.gemma4.target_prefill_contract import ( + file_identity as target_prefill_file_identity, + reviewed_producer_source_path, + validate_target_prefill_receipt, +) SOURCE_CONFIG_SHA256 = ( @@ -95,6 +110,28 @@ }, }, } +ASSISTANT_MODEL_CONTRACT: dict[str, object] = { + "architecture": "Gemma4AssistantForCausalLM", + "backboneHiddenSize": 1536, + "hiddenSize": 256, + "modelType": "gemma4_assistant", + "numHiddenLayers": 4, + "vocabSize": 262144, +} +ASSISTANT_CHECKPOINT_ACQUISITION: dict[str, object] = { + "repo_id": "google/gemma-4-E2B-it-qat-q4_0-unquantized-assistant", + "revision": "ebc7e1a211354561464cb82ed6d886792138dcb6", + "files": { + "config.json": { + "bytes": 2356, + "sha256": "5d01e9f3f8e969aa8147201a26e849c05446c7c746fa918101ed0622b201db15", + }, + "model.safetensors": { + "bytes": 157565344, + "sha256": "28b11aa1fef73e655107984e0024ed1b149df4b8b36dcb95f27cca603eabc960", + }, + }, +} EXPORT_CONTRACT: dict[str, object] = { "backend": "webgpu", "max_input_len": 512, @@ -116,6 +153,9 @@ "[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", ) MTP_EXPORT_CONTRACT: dict[str, object] = { "assistant_calls_per_round": 2, @@ -166,15 +206,187 @@ "legacy_custom_sdpa": 0, "topk": 2, } +_MTP_K2_DONOR_VIEW_ORDER: list[dict[str, object]] = [ + {"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", + }, +] +_MTP_K2_INPUT_ORDER = ["input_ids", "input_pos", "is_round", "donor_length"] +_MTP_K2_OUTPUT_ORDER = [ + "candidates", + "target_greedy", + "output_matches", + "output_bonus", + "state_probe", +] +_MTP_K2_OPERATOR_COUNTS = { + "aten.argmax.default": 3, + "aten.scatter.src": 2, + "aten.topk.default": 2, + "llama.custom_sdpa.default": 43, + "llama.update_cache.default": 31, +} +_MTP_K2_STATE_ALIAS = { + "logicalSource": "nextFeature[1,1,1536]", + "physicalDestination": "seed_feature[1,1,1,1536]", + "mutation": "llama.update_cache.default", +} +_MTP_QAT_DONOR_SEQUENCE = [2, 16, 511, 512, 513, 514, 1024, 8960, 2] +_MTP_QAT_SELECTION_CONTRACT = { + "centroidTopK": 32, + "numCentroids": 2048, + "selectedTokenCount": 4096, + "tokensPerCentroid": 128, +} +COMBINED_RUNTIME_CONTRACT: dict[str, object] = { + "capacities": { + "mtp": { + "donor": 8960, + "max_input": 512, + "target": 8960, + }, + "plain": {"max_context": 8960, "max_input": 512}, + }, + "context": { + "collision": "fail_closed", + "lifetime": "runner_owned", + "registration": "compare_and_set_default_webgpu_context", + "release": "compare_and_set_before_destroy", + }, + "methods": {"mtp": ["k2_round"], "plain": ["text_decoder"]}, + "profile": { + "builds": { + "mtp": { + "profile": "compile_time_enabled", + "wall": "compile_time_disabled", + }, + "plain": {"wall": "compile_time_disabled"}, + }, + "fields": [ + "schemaVersion", + "supported", + "fresh", + "valid", + "context_generation", + "querypool_generation", + "execute_generation", + "total_kernel_ms", + "pass_span_ms", + "interpass_gap_ms", + "perop", + ], + "schema_version": 1, + }, + "reset": { + "mtp": { + "clears": ["accepted_drafts", "buffered_tokens", "execute_count"], + "failure": "fail_closed", + "method": "unload_then_reload", + }, + "plain": { + "failure": "fail_closed", + "method": "unload_then_reload", + }, + }, + "tensor_data_files": {"mtp": 3, "plain": 3}, +} +COMBINED_RUNTIME_VIEWS: dict[str, object] = { + "combined": { + "sequence": ["plain", "mtp", "plain"], + "status": "pending_cross_view_gpu_execution_validation", + }, + "mtp": { + "receipt": "mtp", + "runtime": "mtp", + "status": "pending_gpu_execution_validation", + }, + "plain": { + "receipt": "plain", + "runtime": "plain", + "status": "pending_gpu_execution_validation", + }, +} +RUNTIME_BUILD_TARGETS: dict[str, str] = { + "mtp": "gemma4_spec_browser", + "plain": "gemma4_plain_wasm", +} +RUNTIME_PROFILING_MODES: dict[str, dict[str, bool]] = { + "mtp": {"profile": True, "wall": False}, + "plain": {"profile": True, "wall": False}, +} +RUNTIME_FACTORY_NAMES: dict[str, dict[str, str]] = { + "mtp": { + "profile": "createGemma4MtpProfile", + "wall": "createGemma4Mtp", + }, + "plain": { + "profile": "createWebGPULlama", + "wall": "createWebGPULlama", + }, +} +RUNTIME_OUTPUT_STEMS: dict[str, dict[str, str]] = { + "mtp": { + "profile": "gemma4_mtp_profile", + "wall": "gemma4_mtp", + }, + "plain": { + "profile": "webgpu_llama", + "wall": "webgpu_llama", + }, +} +CLOSURE_SOURCE_PATHS: dict[str, str] = { + "source_manifest": "closure/source_manifest.json", + "wgsl_manifest": "closure/wgsl_manifest.json", +} +CLOSURE_RECIPE_PATHS: dict[str, dict[str, str]] = { + "mtp": { + "profile": "closure/recipes/mtp-profile.json", + "wall": "closure/recipes/mtp-wall.json", + }, + "plain": { + "profile": "closure/recipes/plain-profile.json", + "wall": "closure/recipes/plain-wall.json", + }, +} +RUNTIME_SOURCE_SCHEMA_VERSION = 4 +_BUILD_RECIPE_SCHEMA_VERSION = 2 +_COMBINED_RUNTIME_SCHEMA_VERSION = 3 +_COMMON_WEBGPU_CMAKE_ARGS = [ + "-S", + ".", + "-GNinja", + "-DCMAKE_BUILD_TYPE=Release", + "-DPYTHON_EXECUTABLE=.venv/bin/python", + "-DEXECUTORCH_BUILD_WEBGPU=ON", + "-DEXECUTORCH_BUILD_WEBGPU_TEST=OFF", + "-DEXECUTORCH_BUILD_WASM=ON", + "-DEXECUTORCH_BUILD_XNNPACK=OFF", + "-DEXECUTORCH_BUILD_CPUINFO=ON", + "-DEXECUTORCH_BUILD_PTHREADPOOL=ON", + "-DEXECUTORCH_BUILD_EXTENSION_MODULE=ON", + "-DEXECUTORCH_BUILD_EXTENSION_DATA_LOADER=ON", + "-DEXECUTORCH_BUILD_EXTENSION_TENSOR=ON", + "-DEXECUTORCH_BUILD_EXTENSION_FLAT_TENSOR=ON", + "-DEXECUTORCH_BUILD_EXTENSION_NAMED_DATA_MAP=ON", +] def _source_config_path() -> Path: return Path(__file__).parent / "config" / "e2b_config.json" -def _single_file_manifest( - path: str, byte_count: int, sha256: str -) -> dict[str, object]: +def _single_file_manifest(path: str, byte_count: int, sha256: str) -> dict[str, object]: return { "schema_version": 1, "artifacts": [ @@ -700,6 +912,70 @@ def create_source_closure_receipt( } +def canonical_build_recipe(model: str, flavor: str) -> dict[str, object]: + if model not in RUNTIME_BUILD_TARGETS or flavor not in RUNTIME_PROFILING_MODES.get( + model, {} + ): + raise ValueError(f"unsupported Gemma4 build recipe: {model} {flavor}") + profiling_enabled = RUNTIME_PROFILING_MODES[model][flavor] + build_directory = ( + "cmake-out-gemma4-webgpu-profile" + if profiling_enabled + else "cmake-out-gemma4-webgpu-wall" + ) + output_directory = ( + "backends/webgpu/browser_gemma4_plain" + if model == "plain" + else "backends/webgpu/browser_gemma4_mtp" + ) + output_stem = RUNTIME_OUTPUT_STEMS[model][flavor] + factory = RUNTIME_FACTORY_NAMES[model][flavor] + target = RUNTIME_BUILD_TARGETS[model] + return { + "build_argv": [ + "cmake", + "--build", + build_directory, + "--target", + target, + "-j", + ], + "configure_argv": [ + "emcmake", + "cmake", + *_COMMON_WEBGPU_CMAKE_ARGS, + f"-DEXECUTORCH_BUILD_WEBGPU_PROFILING={'ON' if profiling_enabled else 'OFF'}", + f"-DGEMMA4_SPEC_WASM_EXPORT_NAME={RUNTIME_FACTORY_NAMES['mtp'][flavor]}", + f"-DGEMMA4_SPEC_WASM_OUTPUT_NAME={RUNTIME_OUTPUT_STEMS['mtp'][flavor]}", + "-B", + build_directory, + ], + "cwd": ".", + "factory": factory, + "flavor": flavor, + "model": model, + "output_stem": output_stem, + "outputs": { + "javascript": f"{build_directory}/{output_directory}/{output_stem}.js", + "wasm": f"{build_directory}/{output_directory}/{output_stem}.wasm", + }, + "profiling_enabled": profiling_enabled, + "schema_version": _BUILD_RECIPE_SCHEMA_VERSION, + "target": target, + } + + +def _validate_build_recipe(path: Path, model: str, flavor: str) -> Mapping[str, object]: + label = f"{'MTP' if model == 'mtp' else 'plain'} {flavor} recipe" + try: + recipe = _load_json(path) + except (OSError, json.JSONDecodeError, ValueError) as error: + raise ValueError(f"Gemma4 {label} is not canonical JSON") from error + if recipe != canonical_build_recipe(model, flavor): + raise ValueError(f"Gemma4 {label} does not match the canonical contract") + return recipe + + def _validate_source_receipt(root: Path, artifacts: Sequence[object]) -> None: source_paths = [ artifact.get("path") @@ -751,6 +1027,22 @@ def _validate_source_receipt(root: Path, artifacts: Sequence[object]) -> None: raise ValueError("Gemma4 source receipt verification is incomplete") +def _model_source_receipt( + root: Path, manifest: Mapping[str, object], label: str +) -> Mapping[str, object]: + artifacts = manifest.get("artifacts") + if not isinstance(artifacts, list): + raise ValueError(f"Gemma4 {label} manifest has no artifacts") + source_paths = [ + artifact.get("path") + for artifact in artifacts + if isinstance(artifact, dict) and artifact.get("role") == "source" + ] + if len(source_paths) != 1 or not isinstance(source_paths[0], str): + raise ValueError(f"Gemma4 {label} manifest requires one source receipt") + return _load_json(root / source_paths[0]) + + def validate_export_identity(checkpoint_root: Path) -> Mapping[str, object]: source_config = _source_config_path() validate_manifest( @@ -780,6 +1072,42 @@ def validate_export_identity(checkpoint_root: Path) -> Mapping[str, object]: return CHECKPOINT_ACQUISITION +def validate_assistant_export_identity( + checkpoint_root: Path, +) -> Mapping[str, object]: + files = ASSISTANT_CHECKPOINT_ACQUISITION["files"] + assert isinstance(files, dict) + for name, identity in files.items(): + assert isinstance(name, str) + assert isinstance(identity, dict) + validate_manifest( + checkpoint_root, + _single_file_manifest( + name, + int(identity["bytes"]), + str(identity["sha256"]), + ), + ) + config = _load_json(checkpoint_root / "config.json") + text_config = config.get("text_config") + if not isinstance(text_config, dict): + raise ValueError("assistant config is missing text_config") + architectures = config.get("architectures") + if not isinstance(architectures, list) or not architectures: + raise ValueError("assistant config is missing architectures") + observed_contract = { + "architecture": architectures[0], + "backboneHiddenSize": config.get("backbone_hidden_size"), + "hiddenSize": text_config.get("hidden_size"), + "modelType": config.get("model_type"), + "numHiddenLayers": text_config.get("num_hidden_layers"), + "vocabSize": text_config.get("vocab_size"), + } + if observed_contract != ASSISTANT_MODEL_CONTRACT: + raise ValueError("assistant checkpoint model contract mismatch") + return ASSISTANT_CHECKPOINT_ACQUISITION + + def create_plain_manifest( root: Path, role_paths: Mapping[str, Path], @@ -816,16 +1144,89 @@ def create_plain_manifest( return manifest +def _require_unique_mtp_artifact_paths(artifacts: Sequence[object]) -> None: + seen: set[str] = set() + for artifact in artifacts: + if not isinstance(artifact, dict) or not isinstance(artifact.get("path"), str): + continue + normalized = os.path.normpath(str(artifact["path"])) + if normalized in seen: + raise ValueError( + f"Gemma4 MTP duplicate normalized artifact path: {normalized}" + ) + seen.add(normalized) + + +def create_mtp_manifest( + root: Path, + role_paths: Mapping[str, Path], + ptd_paths: Sequence[Path], +) -> dict[str, object]: + if set(role_paths) not in ({"pte"}, {"pte", "source"}): + raise ValueError( + "Gemma4 MTP manifest requires one K=2 PTE role and an optional " + "source receipt" + ) + if len(ptd_paths) != 3: + raise ValueError("Gemma4 MTP manifest requires exactly three ordered PTDs") + _require_unique_mtp_artifact_paths( + [{"path": str(path)} for path in list(role_paths.values()) + list(ptd_paths)] + ) + manifest = create_manifest(root, role_paths, ptd_paths) + artifacts = manifest.get("artifacts") + assert isinstance(artifacts, list) + _require_unique_mtp_artifact_paths(artifacts) + for artifact in artifacts: + if ( + isinstance(artifact, dict) + and artifact.get("role") == "ptd" + and int(artifact["bytes"]) >= WEBGPU_EXTERNAL_CONSTANTS_MAX_DATA_BYTES + ): + raise ValueError("Gemma4 MTP PTD exceeds the WebGPU constant limit") + paths = [ + str(artifact["path"]) for artifact in artifacts if isinstance(artifact, dict) + ] + if len(set(paths)) != len(paths): + raise ValueError("Gemma4 MTP artifact paths must be distinct") + if any(len(Path(path).parts) != 1 for path in paths): + raise ValueError("Gemma4 MTP artifact paths must be flat") + manifest.update( + { + "acquisition": { + "assistant": ASSISTANT_CHECKPOINT_ACQUISITION, + "target": CHECKPOINT_ACQUISITION, + }, + "export": MTP_EXPORT_CONTRACT, + "model": { + "architecture": ARCHITECTURE_FINGERPRINT, + "assistant": ASSISTANT_MODEL_CONTRACT, + "source_config": MTP_SOURCE_CONFIG, + }, + } + ) + if "source" in role_paths: + _validate_source_receipt(root, artifacts) + manifest["provenance"] = dict(MTP_SOURCE_VERIFIED_PROVENANCE) + else: + manifest["provenance"] = dict(MTP_PENDING_SOURCE_PROVENANCE) + return manifest + + def validate_plain_manifest( root: Path, manifest: Mapping[str, object], require_source_receipt: bool = True, ) -> None: validate_manifest(root, manifest) + if require_source_receipt and manifest.get("provenance") is not None: + raise ValueError("Gemma4 plain production provenance must be absent") if manifest.get("acquisition") != CHECKPOINT_ACQUISITION: raise ValueError("Gemma4 checkpoint acquisition identity mismatch") model = manifest.get("model") - if not isinstance(model, dict) or model.get("architecture") != ARCHITECTURE_FINGERPRINT: + if ( + not isinstance(model, dict) + or model.get("architecture") != ARCHITECTURE_FINGERPRINT + ): raise ValueError("Gemma4 architecture identity mismatch") source_config = model.get("source_config") if source_config != { @@ -850,9 +1251,7 @@ def validate_plain_manifest( ): raise ValueError("Gemma4 PTD exceeds the WebGPU external-constant limit") roles = { - artifact.get("role") - for artifact in artifacts - if isinstance(artifact, dict) + artifact.get("role") for artifact in artifacts if isinstance(artifact, dict) } if "pte" not in roles or (require_source_receipt and "source" not in roles): raise ValueError("Gemma4 plain manifest is missing PTE/source receipt roles") @@ -860,9 +1259,7 @@ def validate_plain_manifest( _validate_source_receipt(root, artifacts) expected_paths = { - str(artifact["path"]) - for artifact in artifacts - if isinstance(artifact, dict) + str(artifact["path"]) for artifact in artifacts if isinstance(artifact, dict) } if any(len(Path(path).parts) != 1 for path in expected_paths): raise ValueError("Gemma4 artifact staging directory must be flat") @@ -871,6 +1268,1623 @@ def validate_plain_manifest( raise ValueError("Gemma4 artifact staging contains missing or extra entries") +def _expected_mtp_mutation_order() -> list[dict[str, object]]: + records: 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"): + records.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", + } + ) + return records + + +def _validate_mtp_k2_evidence(k2_abi: object) -> None: + if not isinstance(k2_abi, dict): + raise ValueError("Gemma4 MTP K=2 ABI evidence must be an object") + _require_exact_keys( + k2_abi, + { + "bufferMutationCount", + "donorViewOrder", + "inputOrder", + "mutationOrder", + "operatorCounts", + "outputOrder", + "seedMutationCount", + "stateAlias", + }, + "Gemma4 MTP K=2 ABI evidence", + ) + for key in ("bufferMutationCount", "seedMutationCount"): + value = k2_abi[key] + if not isinstance(value, int) or isinstance(value, bool) or value < 0: + raise ValueError(f"Gemma4 MTP {key} must be a nonnegative integer") + expected_mutations = _expected_mtp_mutation_order() + if ( + k2_abi["bufferMutationCount"] != len(expected_mutations) + or k2_abi["seedMutationCount"] != 1 + or k2_abi["donorViewOrder"] != _MTP_K2_DONOR_VIEW_ORDER + or k2_abi["inputOrder"] != _MTP_K2_INPUT_ORDER + or k2_abi["mutationOrder"] != expected_mutations + or k2_abi["outputOrder"] != _MTP_K2_OUTPUT_ORDER + or k2_abi["stateAlias"] != _MTP_K2_STATE_ALIAS + ): + raise ValueError("Gemma4 MTP K=2 ABI semantic evidence mismatch") + operator_counts = k2_abi["operatorCounts"] + if not isinstance(operator_counts, dict) or not operator_counts: + raise ValueError("Gemma4 MTP operator counts must be a non-empty object") + if any( + not isinstance(key, str) + or not isinstance(value, int) + or isinstance(value, bool) + or value < 0 + for key, value in operator_counts.items() + ): + raise ValueError("Gemma4 MTP operator counts must be nonnegative integers") + if operator_counts != _MTP_K2_OPERATOR_COUNTS: + raise ValueError("Gemma4 MTP operator-count evidence mismatch") + + +def _validate_mtp_token_record(record: object, label: str) -> None: + if not isinstance(record, dict): + raise ValueError(f"Gemma4 MTP {label} token ordering must be an object") + _require_exact_keys( + record, + { + "max", + "min", + "numel", + "permutationExact", + "rawShape", + "sha256", + "shape", + "uniqueCount", + }, + f"Gemma4 MTP {label} token ordering", + ) + expected = { + "max": 262143, + "min": 0, + "numel": 262144, + "permutationExact": True, + "rawShape": [262144], + "shape": [2048, 128], + "uniqueCount": 262144, + } + if any(record.get(key) != value for key, value in expected.items()): + raise ValueError(f"Gemma4 MTP {label} token-ordering proof mismatch") + if not _is_hex_digest(record.get("sha256"), 64): + raise ValueError(f"Gemma4 MTP {label} token-ordering SHA-256 is invalid") + + +def _validate_mtp_qat_case(case: object, index: int, donor_length: int) -> None: + if not isinstance(case, dict): + raise ValueError("Gemma4 MTP QAT case must be an object") + _require_exact_keys( + case, + { + "caseIndex", + "donorLength", + "greedyTokenExact", + "inputSha256", + "outputs", + "topk", + }, + "Gemma4 MTP QAT case", + ) + input_digests = case["inputSha256"] + outputs = case["outputs"] + if ( + case["caseIndex"] != index + or case["donorLength"] != donor_length + or case["greedyTokenExact"] is not True + or not isinstance(input_digests, list) + or not input_digests + or not all(_is_hex_digest(value, 64) for value in input_digests) + or not isinstance(outputs, list) + or len(outputs) != 2 + ): + raise ValueError("Gemma4 MTP QAT case semantic evidence mismatch") + for output, expected_name in zip(outputs, ("logits", "last_hidden_state")): + if not isinstance(output, dict): + raise ValueError("Gemma4 MTP QAT output evidence must be an object") + _require_exact_keys( + output, + { + "actualSha256", + "bitExact", + "close", + "maxAbsError", + "name", + "referenceSha256", + "shape", + }, + "Gemma4 MTP QAT output evidence", + ) + max_error = output["maxAbsError"] + shape = output["shape"] + valid = ( + output["name"] == expected_name + and isinstance(output["bitExact"], bool) + and output["close"] is True + and isinstance(max_error, (int, float)) + and not isinstance(max_error, bool) + and math.isfinite(float(max_error)) + and float(max_error) >= 0.0 + and isinstance(shape, list) + and bool(shape) + and all( + isinstance(value, int) and not isinstance(value, bool) and value > 0 + for value in shape + ) + and _is_hex_digest(output["actualSha256"], 64) + and _is_hex_digest(output["referenceSha256"], 64) + ) + if not valid: + raise ValueError("Gemma4 MTP QAT output semantic evidence mismatch") + topk = case["topk"] + if not isinstance(topk, dict): + raise ValueError("Gemma4 MTP QAT top-k evidence must be an object") + _require_exact_keys( + topk, + { + "allFinite", + "boundaryGap", + "indicesSha256", + "stableReferenceExact", + "top32PairwiseDistinct", + "top33IndicesSha256", + "top33ValuesSha256", + "valuesSha256", + }, + "Gemma4 MTP QAT top-k evidence", + ) + boundary_gap = topk["boundaryGap"] + valid_topk = ( + topk["allFinite"] is True + and topk["stableReferenceExact"] is True + and topk["top32PairwiseDistinct"] is True + and isinstance(boundary_gap, (int, float)) + and not isinstance(boundary_gap, bool) + and math.isfinite(float(boundary_gap)) + and float(boundary_gap) > 0.0 + and all( + _is_hex_digest(topk[key], 64) + for key in ( + "indicesSha256", + "top33IndicesSha256", + "top33ValuesSha256", + "valuesSha256", + ) + ) + ) + if not valid_topk: + raise ValueError("Gemma4 MTP QAT top-k semantic evidence mismatch") + + +def _validate_mtp_qat_evidence(qat: object) -> None: + if not isinstance(qat, dict): + raise ValueError("Gemma4 MTP QAT evidence must be an object") + _require_exact_keys( + qat, + { + "cases", + "donorSequence", + "eagerEquivalence", + "selectionContract", + "tokenOrdering", + }, + "Gemma4 MTP QAT evidence", + ) + if qat["selectionContract"] != _MTP_QAT_SELECTION_CONTRACT: + raise ValueError("Gemma4 MTP QAT selection contract mismatch") + if qat["eagerEquivalence"] != { + "allClose": True, + "atol": 1e-4, + "rtol": 1e-3, + }: + raise ValueError("Gemma4 MTP eager-equivalence evidence mismatch") + donor_sequence = qat["donorSequence"] + cases = qat["cases"] + if donor_sequence != _MTP_QAT_DONOR_SEQUENCE or not isinstance(cases, list): + raise ValueError("Gemma4 MTP QAT donor sequence mismatch") + if len(cases) != len(donor_sequence): + raise ValueError("Gemma4 MTP QAT case count mismatch") + for index, (case, donor_length) in enumerate(zip(cases, donor_sequence)): + _validate_mtp_qat_case(case, index, donor_length) + if ( + cases[0]["inputSha256"] != cases[-1]["inputSha256"] + or cases[0]["outputs"] != cases[-1]["outputs"] + or cases[0]["topk"] != cases[-1]["topk"] + ): + raise ValueError("Gemma4 MTP QAT replay evidence mismatch") + + token_ordering = qat["tokenOrdering"] + if not isinstance(token_ordering, dict): + raise ValueError("Gemma4 MTP token-ordering evidence must be an object") + base_token_keys = { + "max", + "min", + "numel", + "permutationExact", + "rawShape", + "sha256", + "shape", + "uniqueCount", + } + _require_exact_keys( + token_ordering, + base_token_keys | {"loaded", "raw", "rawLoadedByteExact", "rawSha256"}, + "Gemma4 MTP token-ordering evidence", + ) + effective = {key: token_ordering[key] for key in base_token_keys} + _validate_mtp_token_record(effective, "effective") + _validate_mtp_token_record(token_ordering["raw"], "raw") + _validate_mtp_token_record(token_ordering["loaded"], "loaded") + if ( + token_ordering["rawLoadedByteExact"] is not True + or not _is_hex_digest(token_ordering["rawSha256"], 64) + or token_ordering["raw"] != token_ordering["loaded"] + or effective != token_ordering["loaded"] + or token_ordering["rawSha256"] != token_ordering["sha256"] + ): + raise ValueError("Gemma4 MTP raw/loaded token-ordering identity mismatch") + + +def _validate_mtp_evidence(evidence: object) -> None: # noqa: C901 + if not isinstance(evidence, dict): + raise ValueError("Gemma4 MTP evidence must be an object") + _require_exact_keys( + evidence, + { + "assistant_checkpoint", + "k2_abi", + "lowering", + "qat_selection", + "target_checkpoint", + }, + "Gemma4 MTP evidence", + ) + if evidence["assistant_checkpoint"] != ASSISTANT_CHECKPOINT_ACQUISITION: + raise ValueError("Gemma4 MTP assistant evidence mismatch") + if evidence["target_checkpoint"] != CHECKPOINT_ACQUISITION: + raise ValueError("Gemma4 MTP target evidence mismatch") + _validate_mtp_k2_evidence(evidence["k2_abi"]) + _validate_mtp_qat_evidence(evidence["qat_selection"]) + lowering = evidence["lowering"] + if not isinstance(lowering, dict): + raise ValueError("Gemma4 MTP lowering evidence must be an object") + if lowering != { + "delegate_count": 1, + "edge": MTP_EDGE_CENSUS, + "portable_operator_count": 0, + }: + raise ValueError("Gemma4 MTP lowering census mismatch") + + +def validate_mtp_manifest(root: Path, manifest: Mapping[str, object]) -> None: + provenance = manifest.get("provenance") + requires_source_receipt = False + if provenance == MTP_ACCEPTED_PROVENANCE: + expected_top_level = { + "acquisition", + "artifacts", + "export", + "model", + "provenance", + "ptd_order", + "schema_version", + } + elif provenance in ( + MTP_PENDING_SOURCE_PROVENANCE, + MTP_SOURCE_VERIFIED_PROVENANCE, + ): + expected_top_level = { + "acquisition", + "artifacts", + "evidence", + "export", + "model", + "provenance", + "ptd_order", + "schema_version", + } + _validate_mtp_evidence(manifest.get("evidence")) + requires_source_receipt = provenance == MTP_SOURCE_VERIFIED_PROVENANCE + else: + raise ValueError("Gemma4 MTP provenance mismatch") + _require_exact_keys(manifest, expected_top_level, "Gemma4 MTP manifest") + artifacts = manifest.get("artifacts") + if not isinstance(artifacts, list): + raise ValueError("Gemma4 MTP manifest artifacts must be a list") + _require_unique_mtp_artifact_paths(artifacts) + validate_manifest(root, manifest) + if manifest.get("acquisition") != { + "assistant": ASSISTANT_CHECKPOINT_ACQUISITION, + "target": CHECKPOINT_ACQUISITION, + }: + raise ValueError("Gemma4 MTP checkpoint acquisition identity mismatch") + if manifest.get("model") != { + "architecture": ARCHITECTURE_FINGERPRINT, + "assistant": ASSISTANT_MODEL_CONTRACT, + "source_config": MTP_SOURCE_CONFIG, + }: + raise ValueError("Gemma4 MTP model/source-config identity mismatch") + if manifest.get("export") != MTP_EXPORT_CONTRACT: + raise ValueError("Gemma4 MTP export contract mismatch") + ptd_order = manifest.get("ptd_order") + assert isinstance(ptd_order, list) + if len(ptd_order) != 3: + raise ValueError("Gemma4 MTP manifest requires exactly three ordered PTDs") + roles = [ + artifact.get("role") for artifact in artifacts if isinstance(artifact, dict) + ] + expected_source_count = 1 if requires_source_receipt else 0 + if ( + roles.count("pte") != 1 + or roles.count("ptd") != 3 + or roles.count("source") != expected_source_count + or len(roles) != 4 + expected_source_count + ): + raise ValueError("Gemma4 MTP manifest has unexpected artifact roles") + if requires_source_receipt: + _validate_source_receipt(root, artifacts) + for artifact in artifacts: + assert isinstance(artifact, dict) + _require_exact_keys( + artifact, + {"bytes", "path", "role", "sha256"}, + "Gemma4 MTP artifact", + ) + if ( + artifact.get("role") == "ptd" + and int(artifact["bytes"]) >= WEBGPU_EXTERNAL_CONSTANTS_MAX_DATA_BYTES + ): + raise ValueError("Gemma4 MTP PTD exceeds the WebGPU constant limit") + expected_paths = { + str(artifact["path"]) for artifact in artifacts if isinstance(artifact, dict) + } + if len(expected_paths) != len(artifacts): + raise ValueError("Gemma4 MTP artifact paths must be distinct") + if any(len(Path(path).parts) != 1 for path in expected_paths): + raise ValueError("Gemma4 MTP artifact staging directory must be flat") + actual_paths = {entry.name for entry in root.iterdir()} + if actual_paths != expected_paths: + raise ValueError( + "Gemma4 MTP artifact staging contains missing or extra entries" + ) + + +def _reject_pending_provenance(provenance: object, label: str) -> None: + if not isinstance(provenance, dict): + raise ValueError(f"Gemma4 {label} provenance must be an object") + closure = provenance.get("source_closure") + if not isinstance(closure, str): + raise ValueError(f"Gemma4 {label} provenance has no source closure") + if closure in MTP_PENDING_SOURCE_CLOSURES or closure.startswith("pending"): + raise ValueError( + f"Gemma4 {label} source closure is still pending ({closure}); a " + "pending manifest is never source complete" + ) + + +def _validate_source_complete_mtp_manifest( + root: Path, manifest: Mapping[str, object] +) -> None: + validate_mtp_manifest(root, manifest) + _reject_pending_provenance(manifest.get("provenance"), "MTP manifest") + if manifest.get("provenance") != MTP_SOURCE_VERIFIED_PROVENANCE: + raise ValueError( + "Gemma4 combined runtime requires a source-verified MTP manifest" + ) + artifacts = manifest.get("artifacts") + assert isinstance(artifacts, list) + if not any( + isinstance(artifact, dict) and artifact.get("role") == "source" + for artifact in artifacts + ): + raise ValueError( + "Gemma4 source-verified MTP manifest requires a hashed source receipt" + ) + + +def _contained_regular_file(root: Path, path: Path) -> tuple[Path, str]: + resolved_root = root.resolve(strict=True) + candidate = path if path.is_absolute() else root / path + if candidate.is_symlink(): + raise ValueError(f"runtime staging rejects symlink: {path}") + resolved = candidate.resolve(strict=True) + try: + relative = resolved.relative_to(resolved_root) + except ValueError as error: + raise ValueError(f"runtime artifact escapes staging root: {path}") from error + if not resolved.is_file(): + raise ValueError(f"runtime artifact is not a regular file: {path}") + return resolved, relative.as_posix() + + +def _file_identity(root: Path, path: Path) -> dict[str, object]: + resolved, relative = _contained_regular_file(root, path) + return { + "bytes": resolved.stat().st_size, + "path": relative, + "sha256": _sha256(resolved), + } + + +def _validate_file_identity(root: Path, identity: object, label: str) -> None: + if not isinstance(identity, dict): + raise ValueError(f"{label} identity must be an object") + _require_exact_keys(identity, {"bytes", "path", "sha256"}, label) + path = identity["path"] + if not isinstance(path, str) or Path(path).is_absolute(): + raise ValueError(f"{label} path must be relative") + if not isinstance(identity["bytes"], int) or identity["bytes"] <= 0: + raise ValueError(f"{label} must not be empty") + if not _is_hex_digest(identity["sha256"], 64): + raise ValueError(f"{label} has an invalid SHA-256") + observed = _file_identity(root, Path(path)) + if observed != identity: + raise ValueError(f"{label} byte or SHA-256 identity mismatch") + + +def _runtime_artifact_identity(path: Path) -> dict[str, object]: + if path.is_symlink() or not path.is_file(): + raise ValueError(f"runtime artifact is not a regular file: {path}") + byte_count = path.stat().st_size + if byte_count <= 0: + raise ValueError(f"runtime artifact must not be empty: {path}") + return {"bytes": byte_count, "sha256": _sha256(path)} + + +def _validate_target_prefill_binding( + target_prefill_receipt_path: Path, + runtime_source_receipt_path: Path, +) -> None: + runtime_source_receipt = _load_json(runtime_source_receipt_path) + fbsource_commit = runtime_source_receipt.get("fbsource_commit") + if not _is_hex_digest(fbsource_commit, 40): + raise ValueError("Gemma4 runtime source receipt has an invalid fbsource commit") + producer_path = reviewed_producer_source_path() + producer_identity = target_prefill_file_identity(producer_path) + validate_target_prefill_receipt( + _load_json(target_prefill_receipt_path), + expected_checkpoint_acquisition=CHECKPOINT_ACQUISITION, + expected_producer_path=producer_path, + expected_producer_sha256=str(producer_identity["sha256"]), + expected_runtime_source_identity=_runtime_artifact_identity( + runtime_source_receipt_path + ), + expected_fbsource_commit=str(fbsource_commit), + ) + + +def _closure_identity(path: Path, relative_path: str) -> dict[str, object]: + return {"path": relative_path, **_runtime_artifact_identity(path)} + + +def _identity_bytes_and_hash(identity: Mapping[str, object]) -> tuple[object, object]: + return identity.get("bytes"), identity.get("sha256") + + +def _runtime_role_identity( + build: Mapping[str, object], model: str, role: str +) -> Mapping[str, object]: + identity = build.get(role) + if not isinstance(identity, dict): + raise ValueError(f"Gemma4 {model} {role} identity must be an object") + return identity + + +def _reject_aliased_runtime_roles(runtime: Mapping[str, object]) -> None: + for model in ("mtp", "plain"): + builds = runtime.get(model) + if not isinstance(builds, dict): + raise ValueError(f"Gemma4 {model} runtime build receipt must be an object") + wall = builds.get("wall") + profile = builds.get("profile") + if not isinstance(wall, dict) or not isinstance(profile, dict): + raise ValueError( + f"Gemma4 {model} wall/profile build receipts are incomplete" + ) + wall_javascript = _runtime_role_identity(wall, model, "javascript") + profile_javascript = _runtime_role_identity(profile, model, "javascript") + wall_wasm = _runtime_role_identity(wall, model, "wasm") + profile_wasm = _runtime_role_identity(profile, model, "wasm") + wall_pair = ( + _identity_bytes_and_hash(wall_javascript), + _identity_bytes_and_hash(wall_wasm), + ) + profile_pair = ( + _identity_bytes_and_hash(profile_javascript), + _identity_bytes_and_hash(profile_wasm), + ) + if wall_pair == profile_pair: + raise ValueError(f"Gemma4 {model} wall/profile pair identities must differ") + if _identity_bytes_and_hash(wall_wasm) == _identity_bytes_and_hash( + profile_wasm + ): + raise ValueError(f"Gemma4 {model} wall/profile wasm identities must differ") + for role in ("recipe",): + wall_identity = _runtime_role_identity(wall, model, role) + profile_identity = _runtime_role_identity(profile, model, role) + if _identity_bytes_and_hash(wall_identity) == _identity_bytes_and_hash( + profile_identity + ): + raise ValueError( + f"Gemma4 {model} wall/profile {role} identities must differ" + ) + if model == "mtp" and _identity_bytes_and_hash( + wall_javascript + ) == _identity_bytes_and_hash(profile_javascript): + raise ValueError( + "Gemma4 MTP wall/profile javascript identities must differ" + ) + + +def _validate_runtime_product_basenames( + runtime_paths: Mapping[str, Mapping[str, Mapping[str, Path]]], +) -> None: + for model, builds in runtime_paths.items(): + for flavor, paths in builds.items(): + output_stem = RUNTIME_OUTPUT_STEMS[model][flavor] + for kind, suffix in (("javascript", "js"), ("wasm", "wasm")): + if paths[kind].name != f"{output_stem}.{suffix}": + display_kind = "JavaScript" if kind == "javascript" else "WASM" + raise ValueError( + f"Gemma4 {model} {flavor} {display_kind} basename mismatch" + ) + + +def _validate_runtime_source_receipt( + receipt: Mapping[str, object], + runtime_paths: Mapping[str, Mapping[str, Mapping[str, Path]]], + manifest_paths: Mapping[str, Path], + model_roots: Mapping[str, Path], + source_input_paths: Mapping[str, Path], + build_recipe_paths: Mapping[str, Mapping[str, Path]], +) -> None: + _require_exact_keys( + receipt, + { + "fbsource_commit", + "model_manifests", + "oss_commit", + "runtime", + "schema_version", + "source_inputs", + "source_current", + "verification", + }, + "Gemma4 runtime source receipt", + ) + if receipt.get("schema_version") != RUNTIME_SOURCE_SCHEMA_VERSION: + raise ValueError("Gemma4 runtime source receipt requires schema version 4") + if receipt.get("source_current") is not True: + raise ValueError("Gemma4 runtime source receipt is not source-current") + if receipt.get("verification") != { + "build_execution": "not_attested", + "recipe": "validated", + "source_checkout": "verified", + "wgsl_codegen": "verified", + }: + raise ValueError("Gemma4 runtime source verification claims are invalid") + source_inputs = receipt["source_inputs"] + if not isinstance(source_inputs, dict): + raise ValueError("Gemma4 runtime source inputs must be an object") + _require_exact_keys( + source_inputs, set(CLOSURE_SOURCE_PATHS), "Gemma4 runtime source inputs" + ) + if set(source_input_paths) != set(CLOSURE_SOURCE_PATHS): + raise ValueError("Gemma4 runtime source-input paths are incomplete") + for label, relative_path in CLOSURE_SOURCE_PATHS.items(): + expected = source_inputs[label] + if not isinstance(expected, dict): + raise ValueError(f"Gemma4 {label} identity must be an object") + _require_exact_keys( + expected, {"bytes", "path", "sha256"}, f"Gemma4 {label} identity" + ) + if expected != _closure_identity(source_input_paths[label], relative_path): + raise ValueError(f"Gemma4 {label} is not bound to its source receipt") + source_manifest = _load_json(source_input_paths["source_manifest"]) + wgsl_manifest = _load_json(source_input_paths["wgsl_manifest"]) + try: + validate_source_manifest(source_manifest) + except ValueError as error: + raise ValueError(f"Gemma4 source manifest is invalid: {error}") from error + try: + validate_wgsl_manifest(wgsl_manifest) + except ValueError as error: + raise ValueError(f"Gemma4 WGSL manifest is invalid: {error}") from error + checkouts = source_manifest["checkouts"] + assert isinstance(checkouts, dict) + fbsource_checkout = checkouts["fbsource"] + assert isinstance(fbsource_checkout, dict) + if wgsl_manifest.get("fbsource_commit") != fbsource_checkout.get("head"): + raise ValueError("Gemma4 source and WGSL manifests have different heads") + for label, key in (("fbsource", "fbsource_commit"), ("oss", "oss_commit")): + checkout = checkouts[label] + assert isinstance(checkout, dict) + if receipt.get(key) != checkout.get("head"): + raise ValueError(f"Gemma4 runtime source receipt has invalid {key}") + model_manifests = receipt["model_manifests"] + if not isinstance(model_manifests, dict): + raise ValueError("Gemma4 model-manifest bindings must be an object") + _require_exact_keys( + model_manifests, {"mtp", "plain"}, "Gemma4 model-manifest bindings" + ) + if set(manifest_paths) != {"mtp", "plain"}: + raise ValueError("Gemma4 model-manifest paths are incomplete") + if set(model_roots) != {"mtp", "plain"}: + raise ValueError("Gemma4 model artifact roots are incomplete") + for label in ("plain", "mtp"): + expected = model_manifests[label] + if not isinstance(expected, dict): + raise ValueError(f"Gemma4 {label} manifest identity must be an object") + _require_exact_keys( + expected, {"bytes", "sha256"}, f"Gemma4 {label} manifest identity" + ) + if expected != _runtime_artifact_identity(manifest_paths[label]): + raise ValueError( + f"Gemma4 {label} manifest is not bound to its source receipt" + ) + manifest = _load_json(manifest_paths[label]) + if label == "plain": + validate_plain_manifest(model_roots[label], manifest) + else: + _validate_source_complete_mtp_manifest(model_roots[label], manifest) + model_source = _model_source_receipt(model_roots[label], manifest, label) + if model_source.get("fbsource_commit") != receipt.get( + "fbsource_commit" + ) or model_source.get("oss_commit") != receipt.get("oss_commit"): + raise ValueError(f"Gemma4 {label} source receipt checkout head mismatch") + if ( + model_source.get("source_manifest") != source_manifest + or model_source.get("wgsl_manifest") != wgsl_manifest + ): + raise ValueError(f"Gemma4 {label} source receipt closure mismatch") + runtime = receipt["runtime"] + if not isinstance(runtime, dict): + raise ValueError("Gemma4 runtime source receipt runtime must be an object") + _require_exact_keys(runtime, set(RUNTIME_BUILD_TARGETS), "Gemma4 runtime builds") + if set(runtime_paths) != set(RUNTIME_BUILD_TARGETS): + raise ValueError("Gemma4 runtime paths are incomplete") + if set(build_recipe_paths) != set(RUNTIME_BUILD_TARGETS): + raise ValueError("Gemma4 build-recipe paths are incomplete") + for model, target in RUNTIME_BUILD_TARGETS.items(): + builds = runtime[model] + if not isinstance(builds, dict): + raise ValueError(f"Gemma4 {model} build receipt must be an object") + flavors = RUNTIME_PROFILING_MODES[model] + _require_exact_keys( + builds, + {"target", *flavors}, + f"Gemma4 {model} build receipt", + ) + if builds.get("target") != target: + raise ValueError(f"Gemma4 {model} runtime build target mismatch") + model_paths = runtime_paths[model] + model_recipe_paths = build_recipe_paths[model] + if set(model_paths) != set(flavors): + raise ValueError(f"Gemma4 {model} runtime paths are incomplete") + if set(model_recipe_paths) != set(flavors): + raise ValueError(f"Gemma4 {model} build-recipe paths are incomplete") + for flavor, profiling_enabled in flavors.items(): + build = builds[flavor] + if not isinstance(build, dict): + raise ValueError( + f"Gemma4 {model} {flavor} build receipt must be an object" + ) + _require_exact_keys( + build, + { + "factory", + "javascript", + "output_stem", + "profiling_enabled", + "recipe", + "wasm", + }, + f"Gemma4 {model} {flavor} build receipt", + ) + if build["profiling_enabled"] is not profiling_enabled: + raise ValueError(f"Gemma4 {model} {flavor} profiling mode mismatch") + recipe = build["recipe"] + if not isinstance(recipe, dict): + raise ValueError( + f"Gemma4 {model} {flavor} recipe identity must be an object" + ) + _require_exact_keys( + recipe, + {"bytes", "path", "sha256"}, + f"Gemma4 {model} {flavor} recipe identity", + ) + if recipe != _closure_identity( + model_recipe_paths[flavor], CLOSURE_RECIPE_PATHS[model][flavor] + ): + raise ValueError( + f"Gemma4 {model} {flavor} recipe is not bound to its build receipt" + ) + recipe_document = _validate_build_recipe( + model_recipe_paths[flavor], model, flavor + ) + if build["factory"] != recipe_document["factory"]: + raise ValueError(f"Gemma4 {model} {flavor} factory mismatch") + if build["output_stem"] != recipe_document["output_stem"]: + raise ValueError(f"Gemma4 {model} {flavor} output_stem mismatch") + paths = model_paths[flavor] + if set(paths) != {"javascript", "wasm"}: + raise ValueError( + f"Gemma4 {model} {flavor} runtime paths are incomplete" + ) + for kind in ("javascript", "wasm"): + expected = build[kind] + if not isinstance(expected, dict): + raise ValueError( + f"Gemma4 {model} {flavor} {kind} identity must be an object" + ) + _require_exact_keys( + expected, + {"bytes", "sha256"}, + f"Gemma4 {model} {flavor} {kind} identity", + ) + if expected != _runtime_artifact_identity(paths[kind]): + raise ValueError( + f"Gemma4 {model} {flavor} {kind} is not bound to its build receipt" + ) + _reject_aliased_runtime_roles(runtime) + + +def _runtime_paths(root: Path) -> dict[str, dict[str, dict[str, Path]]]: + return { + model: { + flavor: { + "javascript": root / f"runtime/{model}/{flavor}.js", + "wasm": root / f"runtime/{model}/{flavor}.wasm", + } + for flavor in flavors + } + for model, flavors in { + "mtp": ("profile", "wall"), + "plain": ("profile", "wall"), + }.items() + } + + +def _closure_source_paths(root: Path) -> dict[str, Path]: + return {label: root / path for label, path in CLOSURE_SOURCE_PATHS.items()} + + +def _closure_recipe_paths(root: Path) -> dict[str, dict[str, Path]]: + return { + model: {flavor: root / path for flavor, path in recipes.items()} + for model, recipes in CLOSURE_RECIPE_PATHS.items() + } + + +def create_runtime_source_receipt( + *, + fbsource_root: Path, + oss_root: Path, + backend_root: Path, + source_manifest_path: Path, + wgsl_manifest_path: Path, + manifest_paths: Mapping[str, Path], + model_roots: Mapping[str, Path], + runtime_paths: Mapping[str, Mapping[str, Mapping[str, Path]]], + build_command_paths: Mapping[str, Mapping[str, Path]], +) -> dict[str, object]: + source_manifest = _load_json(source_manifest_path) + wgsl_manifest = _load_json(wgsl_manifest_path) + try: + validate_source_manifest(source_manifest) + except ValueError as error: + raise ValueError(f"Gemma4 source manifest is invalid: {error}") from error + try: + validate_wgsl_manifest(wgsl_manifest) + except ValueError as error: + raise ValueError(f"Gemma4 WGSL manifest is invalid: {error}") from error + live_source_manifest = create_source_manifest(fbsource_root, oss_root) + if source_manifest != live_source_manifest: + raise ValueError( + "Gemma4 source manifest does not match the live clean checkouts" + ) + live_wgsl_manifest = create_wgsl_manifest(backend_root) + if wgsl_manifest != live_wgsl_manifest: + raise ValueError( + "Gemma4 WGSL manifest does not match the live generator closure" + ) + checkouts = source_manifest["checkouts"] + assert isinstance(checkouts, dict) + fbsource_checkout = checkouts["fbsource"] + oss_checkout = checkouts["oss"] + assert isinstance(fbsource_checkout, dict) and isinstance(oss_checkout, dict) + fbsource_commit = str(fbsource_checkout["head"]) + oss_commit = str(oss_checkout["head"]) + _require_exact_keys(manifest_paths, {"mtp", "plain"}, "Gemma4 model-manifest paths") + _require_exact_keys(model_roots, {"mtp", "plain"}, "Gemma4 model artifact roots") + _require_exact_keys( + runtime_paths, set(RUNTIME_BUILD_TARGETS), "Gemma4 runtime paths" + ) + _require_exact_keys( + build_command_paths, + set(RUNTIME_BUILD_TARGETS), + "Gemma4 build-command paths", + ) + for model, recipes in build_command_paths.items(): + if _identity_bytes_and_hash( + _runtime_artifact_identity(recipes["wall"]) + ) == _identity_bytes_and_hash(_runtime_artifact_identity(recipes["profile"])): + raise ValueError( + f"Gemma4 {model} wall/profile recipe identities must differ" + ) + + runtime: dict[str, object] = {} + for model, target in RUNTIME_BUILD_TARGETS.items(): + expected_flavors = RUNTIME_PROFILING_MODES[model] + model_runtime_paths = runtime_paths[model] + model_command_paths = build_command_paths[model] + _require_exact_keys( + model_runtime_paths, + set(expected_flavors), + f"Gemma4 {model} runtime paths", + ) + _require_exact_keys( + model_command_paths, + set(expected_flavors), + f"Gemma4 {model} build-command paths", + ) + builds: dict[str, object] = {"target": target} + for flavor, profiling_enabled in expected_flavors.items(): + _validate_build_recipe(model_command_paths[flavor], model, flavor) + artifacts = model_runtime_paths[flavor] + _require_exact_keys( + artifacts, + {"javascript", "wasm"}, + f"Gemma4 {model} {flavor} runtime paths", + ) + builds[flavor] = { + "factory": RUNTIME_FACTORY_NAMES[model][flavor], + "javascript": _runtime_artifact_identity(artifacts["javascript"]), + "output_stem": RUNTIME_OUTPUT_STEMS[model][flavor], + "profiling_enabled": profiling_enabled, + "recipe": _closure_identity( + model_command_paths[flavor], + CLOSURE_RECIPE_PATHS[model][flavor], + ), + "wasm": _runtime_artifact_identity(artifacts["wasm"]), + } + runtime[model] = builds + + source_input_paths = { + "source_manifest": source_manifest_path, + "wgsl_manifest": wgsl_manifest_path, + } + receipt: dict[str, object] = { + "fbsource_commit": fbsource_commit, + "model_manifests": { + label: _runtime_artifact_identity(path) + for label, path in manifest_paths.items() + }, + "oss_commit": oss_commit, + "runtime": runtime, + "schema_version": RUNTIME_SOURCE_SCHEMA_VERSION, + "source_inputs": { + label: _closure_identity(source_input_paths[label], relative_path) + for label, relative_path in CLOSURE_SOURCE_PATHS.items() + }, + "source_current": True, + "verification": { + "build_execution": "not_attested", + "recipe": "validated", + "source_checkout": "verified", + "wgsl_codegen": "verified", + }, + } + _validate_runtime_source_receipt( + receipt, + runtime_paths, + manifest_paths, + model_roots, + source_input_paths, + build_command_paths, + ) + _validate_runtime_product_basenames(runtime_paths) + return receipt + + +def _staged_regular_files(root: Path) -> set[str]: + files: set[str] = set() + for directory, directories, filenames in os.walk(root): + current = Path(directory) + for name in directories: + path = current / name + if path.is_symlink(): + raise ValueError(f"runtime staging rejects symlink: {path}") + for name in filenames: + path = current / name + if path.is_symlink() or not path.is_file(): + raise ValueError(f"runtime staging rejects non-regular file: {path}") + files.add(path.relative_to(root).as_posix()) + return files + + +def _source_receipt_identity( + root: Path, manifest: Mapping[str, object], label: str +) -> dict[str, object]: + artifacts = manifest.get("artifacts") + if not isinstance(artifacts, list): + raise ValueError(f"Gemma4 {label} manifest has no artifacts") + entries = [ + artifact + for artifact in artifacts + if isinstance(artifact, dict) and artifact.get("role") == "source" + ] + if len(entries) != 1: + raise ValueError(f"Gemma4 {label} manifest requires one source receipt") + entry = entries[0] + return {key: entry[key] for key in ("bytes", "path", "sha256")} + + +def _source_verification( + root: Path, + plain_manifest: Mapping[str, object], + mtp_manifest: Mapping[str, object], +) -> dict[str, object]: + verification: dict[str, object] = {} + for label, manifest in (("mtp", mtp_manifest), ("plain", plain_manifest)): + provenance = manifest.get("provenance") + if label == "mtp": + _reject_pending_provenance(provenance, "MTP manifest") + verification[label] = { + "provenance": copy.deepcopy(provenance) if provenance else None, + "source_receipt": _source_receipt_identity(root / label, manifest, label), + } + return verification + + +def create_combined_runtime_envelope(root: Path) -> dict[str, object]: + plain_receipt_path = Path("receipts/plain.json") + mtp_receipt_path = Path("receipts/mtp.json") + runtime_source_path = Path("receipts/runtime_source.json") + target_prefill_path = Path("receipts/target_prefill.json") + + plain_receipt = _load_json(root / plain_receipt_path) + mtp_receipt = _load_json(root / mtp_receipt_path) + runtime_source_receipt = _load_json(root / runtime_source_path) + validate_plain_manifest(root / "plain", plain_receipt) + _validate_source_complete_mtp_manifest(root / "mtp", mtp_receipt) + _validate_runtime_source_receipt( + runtime_source_receipt, + _runtime_paths(root), + {"mtp": root / mtp_receipt_path, "plain": root / plain_receipt_path}, + {"mtp": root / "mtp", "plain": root / "plain"}, + _closure_source_paths(root), + _closure_recipe_paths(root), + ) + _validate_target_prefill_binding( + root / target_prefill_path, root / runtime_source_path + ) + runtime = { + model: { + flavor: { + kind: _file_identity(root, path) for kind, path in artifacts.items() + } + for flavor, artifacts in builds.items() + } + for model, builds in _runtime_paths(root).items() + } + + envelope: dict[str, object] = { + "contract": COMBINED_RUNTIME_CONTRACT, + "receipts": { + "mtp": { + **_file_identity(root, mtp_receipt_path), + "root": "mtp", + }, + "plain": { + **_file_identity(root, plain_receipt_path), + "root": "plain", + }, + "target_prefill": _file_identity(root, target_prefill_path), + }, + "runtime": {**runtime, "source": _file_identity(root, runtime_source_path)}, + "schema_version": _COMBINED_RUNTIME_SCHEMA_VERSION, + "source_verification": _source_verification(root, plain_receipt, mtp_receipt), + "views": COMBINED_RUNTIME_VIEWS, + } + return envelope + + +def validate_combined_runtime_envelope( + root: Path, envelope: Mapping[str, object] +) -> None: + _require_exact_keys( + envelope, + { + "contract", + "receipts", + "runtime", + "schema_version", + "source_verification", + "views", + }, + "Gemma4 combined runtime envelope", + ) + if envelope.get("schema_version") != _COMBINED_RUNTIME_SCHEMA_VERSION: + raise ValueError("Gemma4 combined runtime schema version mismatch") + if envelope.get("contract") != COMBINED_RUNTIME_CONTRACT: + raise ValueError("Gemma4 combined runtime contract mismatch") + if envelope.get("views") != COMBINED_RUNTIME_VIEWS: + raise ValueError("Gemma4 combined runtime views mismatch") + source_verification = envelope.get("source_verification") + if not isinstance(source_verification, dict): + raise ValueError("Gemma4 combined source verification must be an object") + _require_exact_keys( + source_verification, {"mtp", "plain"}, "Gemma4 combined source verification" + ) + for label in ("mtp", "plain"): + entry = source_verification[label] + if not isinstance(entry, dict): + raise ValueError(f"Gemma4 {label} source verification must be an object") + _require_exact_keys( + entry, + {"provenance", "source_receipt"}, + f"Gemma4 {label} source verification", + ) + receipt = entry["source_receipt"] + if not isinstance(receipt, dict): + raise ValueError(f"Gemma4 {label} source receipt must be an object") + _require_exact_keys( + receipt, {"bytes", "path", "sha256"}, f"Gemma4 {label} source receipt" + ) + _validate_file_identity(root / label, receipt, f"Gemma4 {label} source receipt") + _reject_pending_provenance( + source_verification["mtp"].get("provenance"), "MTP manifest" + ) + if source_verification["mtp"].get("provenance") != MTP_SOURCE_VERIFIED_PROVENANCE: + raise ValueError("Gemma4 combined runtime requires source-verified MTP") + + receipts = envelope.get("receipts") + if not isinstance(receipts, dict): + raise ValueError("Gemma4 combined receipts must be an object") + _require_exact_keys( + receipts, + {"mtp", "plain", "target_prefill"}, + "Gemma4 combined receipts", + ) + expected_paths = { + "receipts/mtp.json", + "receipts/plain.json", + "receipts/runtime_source.json", + "receipts/target_prefill.json", + } + expected_paths.update( + path.relative_to(root).as_posix() + for builds in _runtime_paths(root).values() + for artifacts in builds.values() + for path in artifacts.values() + ) + expected_paths.update(CLOSURE_SOURCE_PATHS.values()) + expected_paths.update( + path for recipes in CLOSURE_RECIPE_PATHS.values() for path in recipes.values() + ) + for label in ("plain", "mtp"): + receipt = receipts[label] + if not isinstance(receipt, dict): + raise ValueError(f"Gemma4 {label} receipt must be an object") + _require_exact_keys( + receipt, {"bytes", "path", "root", "sha256"}, f"Gemma4 {label} receipt" + ) + receipt_identity = {key: receipt[key] for key in ("bytes", "path", "sha256")} + _validate_file_identity(root, receipt_identity, f"Gemma4 {label} receipt") + if receipt.get("root") != label: + raise ValueError(f"Gemma4 {label} receipt root mismatch") + manifest = _load_json(root / str(receipt["path"])) + if label == "plain": + validate_plain_manifest(root / label, manifest) + else: + _validate_source_complete_mtp_manifest(root / label, manifest) + artifacts = manifest.get("artifacts") + assert isinstance(artifacts, list) + expected_paths.update( + f"{label}/{artifact['path']}" + for artifact in artifacts + if isinstance(artifact, dict) + ) + + target_prefill = receipts["target_prefill"] + if not isinstance(target_prefill, dict): + raise ValueError("Gemma4 target-prefill receipt must be an object") + _require_exact_keys( + target_prefill, + {"bytes", "path", "sha256"}, + "Gemma4 target-prefill receipt", + ) + if target_prefill.get("path") != "receipts/target_prefill.json": + raise ValueError("Gemma4 target-prefill receipt path mismatch") + _validate_file_identity(root, target_prefill, "Gemma4 target-prefill receipt") + + runtime = envelope.get("runtime") + if not isinstance(runtime, dict): + raise ValueError("Gemma4 combined runtime must be an object") + _require_exact_keys(runtime, {"mtp", "plain", "source"}, "Gemma4 combined runtime") + _validate_file_identity(root, runtime["source"], "Gemma4 runtime source") + source_identity = runtime["source"] + assert isinstance(source_identity, dict) + _validate_runtime_source_receipt( + _load_json(root / str(source_identity["path"])), + _runtime_paths(root), + { + "mtp": root / "receipts/mtp.json", + "plain": root / "receipts/plain.json", + }, + {"mtp": root / "mtp", "plain": root / "plain"}, + _closure_source_paths(root), + _closure_recipe_paths(root), + ) + _validate_target_prefill_binding( + root / "receipts/target_prefill.json", + root / str(source_identity["path"]), + ) + for model, expected_builds in _runtime_paths(root).items(): + model_runtime = runtime[model] + if not isinstance(model_runtime, dict): + raise ValueError(f"Gemma4 {model} runtime must be an object") + _require_exact_keys( + model_runtime, set(expected_builds), f"Gemma4 {model} runtime" + ) + for flavor in expected_builds: + artifacts = model_runtime[flavor] + if not isinstance(artifacts, dict): + raise ValueError(f"Gemma4 {model} {flavor} runtime must be an object") + _require_exact_keys( + artifacts, + {"javascript", "wasm"}, + f"Gemma4 {model} {flavor} runtime", + ) + _validate_file_identity( + root, + artifacts["javascript"], + f"Gemma4 {model} {flavor} JavaScript", + ) + _validate_file_identity( + root, artifacts["wasm"], f"Gemma4 {model} {flavor} WASM" + ) + extra_paths = ( + _staged_regular_files(root) + - expected_paths + - {"gemma4_webgpu_combined_runtime.json"} + ) + if extra_paths: + raise ValueError( + f"Gemma4 combined runtime contains extra files: {sorted(extra_paths)}" + ) + if envelope != create_combined_runtime_envelope(root): + raise ValueError("Gemma4 combined runtime has non-canonical role bindings") + + +def _copy_manifest_artifacts( + source_root: Path, + manifest: Mapping[str, object], + destination_root: Path, +) -> None: + artifacts = manifest.get("artifacts") + assert isinstance(artifacts, list) + destination_root.mkdir() + for artifact in artifacts: + assert isinstance(artifact, dict) + path = Path(str(artifact["path"])) + source, relative = _contained_regular_file(source_root, path) + destination = destination_root / relative + destination.parent.mkdir(parents=True, exist_ok=True) + shutil.copyfile(source, destination) + + +def stage_combined_runtime( + destination_root: Path, + plain_root: Path, + plain_receipt_path: Path, + mtp_root: Path, + mtp_receipt_path: Path, + runtime_source_receipt_path: Path, + mtp_wall_javascript_path: Path, + mtp_wall_wasm_path: Path, + mtp_profile_javascript_path: Path, + mtp_profile_wasm_path: Path, + *, + plain_profile_javascript_path: Path, + plain_profile_wasm_path: Path, + plain_wall_javascript_path: Path, + plain_wall_wasm_path: Path, + source_manifest_path: Path, + wgsl_manifest_path: Path, + build_recipe_paths: Mapping[str, Mapping[str, Path]], + target_prefill_receipt_path: Path, +) -> None: + if destination_root.exists() or destination_root.is_symlink(): + raise ValueError( + f"runtime staging destination already exists: {destination_root}" + ) + destination_root.parent.mkdir(parents=True, exist_ok=True) + + for receipt_path in ( + plain_receipt_path, + mtp_receipt_path, + runtime_source_receipt_path, + target_prefill_receipt_path, + ): + if receipt_path.is_symlink() or not receipt_path.is_file(): + raise ValueError( + f"runtime staging receipt is not a regular file: {receipt_path}" + ) + plain_receipt = _load_json(plain_receipt_path) + mtp_receipt = _load_json(mtp_receipt_path) + validate_plain_manifest(plain_root, plain_receipt) + _validate_source_complete_mtp_manifest(mtp_root, mtp_receipt) + runtime_inputs = { + "mtp": { + "profile": { + "javascript": mtp_profile_javascript_path, + "wasm": mtp_profile_wasm_path, + }, + "wall": { + "javascript": mtp_wall_javascript_path, + "wasm": mtp_wall_wasm_path, + }, + }, + "plain": { + "profile": { + "javascript": plain_profile_javascript_path, + "wasm": plain_profile_wasm_path, + }, + "wall": { + "javascript": plain_wall_javascript_path, + "wasm": plain_wall_wasm_path, + }, + }, + } + source_inputs = { + "source_manifest": source_manifest_path, + "wgsl_manifest": wgsl_manifest_path, + } + for builds in runtime_inputs.values(): + for artifacts in builds.values(): + for path in artifacts.values(): + if path.is_symlink() or not path.is_file(): + raise ValueError( + f"runtime staging input is not a regular file: {path}" + ) + closure_inputs = list(source_inputs.values()) + for builds in build_recipe_paths.values(): + closure_inputs.extend(builds.values()) + for path in closure_inputs: + if path.is_symlink() or not path.is_file(): + raise ValueError(f"runtime closure input is not a regular file: {path}") + if ( + runtime_source_receipt_path.is_symlink() + or not runtime_source_receipt_path.is_file() + ): + raise ValueError( + "runtime staging source receipt is not a regular file: " + f"{runtime_source_receipt_path}" + ) + _validate_runtime_source_receipt( + _load_json(runtime_source_receipt_path), + runtime_inputs, + {"mtp": mtp_receipt_path, "plain": plain_receipt_path}, + {"mtp": mtp_root, "plain": plain_root}, + source_inputs, + build_recipe_paths, + ) + _validate_target_prefill_binding( + target_prefill_receipt_path, runtime_source_receipt_path + ) + + temporary_root = Path( + tempfile.mkdtemp( + prefix=f".{destination_root.name}.", dir=destination_root.parent + ) + ) + try: + _copy_manifest_artifacts(plain_root, plain_receipt, temporary_root / "plain") + _copy_manifest_artifacts(mtp_root, mtp_receipt, temporary_root / "mtp") + receipts_root = temporary_root / "receipts" + receipts_root.mkdir() + shutil.copyfile(plain_receipt_path, receipts_root / "plain.json") + shutil.copyfile(mtp_receipt_path, receipts_root / "mtp.json") + shutil.copyfile( + runtime_source_receipt_path, receipts_root / "runtime_source.json" + ) + shutil.copyfile( + target_prefill_receipt_path, receipts_root / "target_prefill.json" + ) + for label, destination in CLOSURE_SOURCE_PATHS.items(): + staged_path = temporary_root / destination + staged_path.parent.mkdir(parents=True, exist_ok=True) + shutil.copyfile(source_inputs[label], staged_path) + for model, recipes in CLOSURE_RECIPE_PATHS.items(): + for flavor, destination in recipes.items(): + staged_path = temporary_root / destination + staged_path.parent.mkdir(parents=True, exist_ok=True) + shutil.copyfile(build_recipe_paths[model][flavor], staged_path) + for model, builds in runtime_inputs.items(): + runtime_root = temporary_root / "runtime" / model + runtime_root.mkdir(parents=True) + for flavor, artifacts in builds.items(): + for kind, source in artifacts.items(): + suffix = "js" if kind == "javascript" else "wasm" + shutil.copyfile(source, runtime_root / f"{flavor}.{suffix}") + + envelope = create_combined_runtime_envelope(temporary_root) + validate_combined_runtime_envelope(temporary_root, envelope) + (temporary_root / "gemma4_webgpu_combined_runtime.json").write_text( + json.dumps(envelope, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + os.replace(temporary_root, destination_root) + except BaseException: + shutil.rmtree(temporary_root, ignore_errors=True) + raise + + +@dataclasses.dataclass(slots=True) +class _PublishedArtifact: + destination: Path + expected_device: int + expected_inode: int + owned: bool = False + + +def _path_identity(path: Path) -> tuple[int, int] | None: + try: + observed = path.stat(follow_symlinks=False) + except FileNotFoundError: + return None + return observed.st_dev, observed.st_ino + + +def _quarantine_owned_path( + path: Path, + expected_identity: tuple[int, int], + quarantine_parent: Path, +) -> bool: + quarantine_root = Path( + tempfile.mkdtemp(prefix=".mtp-publication-quarantine.", dir=quarantine_parent) + ) + quarantined = quarantine_root / path.name + try: + try: + os.rename(path, quarantined) + except FileNotFoundError: + return False + + observed_identity = _path_identity(quarantined) + if observed_identity == expected_identity: + quarantined.unlink() + return True + if observed_identity is None: + return False + + try: + os.link(quarantined, path, follow_symlinks=False) + except OSError as error: + raise RuntimeError( + f"foreign publication entry retained for recovery at {quarantined}" + ) from error + if _path_identity(path) != observed_identity: + raise RuntimeError( + f"foreign publication entry retained for recovery at {quarantined}" + ) + quarantined.unlink() + return False + finally: + try: + quarantine_root.rmdir() + except OSError as error: + if error.errno not in (errno.EEXIST, errno.ENOTEMPTY): + raise + + +def _link_and_claim( + staged: Path, + destination: Path, + publication: _PublishedArtifact, +) -> None: + """Defer SIGINT through a normal-return transition to owned. + + SIGKILL or another exception with uncertain syscall completion can leave a + pre-receipt partial file. The final receipt is the commit witness, and a + later no-clobber publication attempt fails closed on that partial. + """ + if threading.current_thread() is not threading.main_thread(): + os.link(staged, destination, follow_symlinks=False) + publication.owned = True + return + + received_sigint = False + + def defer_sigint(_signum: int, _frame: FrameType | None) -> None: + nonlocal received_sigint + received_sigint = True + + previous_handler = signal.getsignal(signal.SIGINT) + signal.signal(signal.SIGINT, defer_sigint) + try: + os.link(staged, destination, follow_symlinks=False) + # Keep this as the only operation between a normal link return and OWNED. + publication.owned = True + finally: + signal.signal(signal.SIGINT, previous_handler) + if received_sigint: + signal.raise_signal(signal.SIGINT) + + +def _publish_no_clobber( + staged: Path, + destination: Path, + published: list[_PublishedArtifact], +) -> None: + source_stat = staged.stat(follow_symlinks=False) + if not stat.S_ISREG(source_stat.st_mode): + raise ValueError(f"staged publication is not regular: {staged}") + if destination.parent.stat(follow_symlinks=False).st_dev != source_stat.st_dev: + raise ValueError("staged and final artifacts must use the same filesystem") + + expected_identity = (source_stat.st_dev, source_stat.st_ino) + publication = _PublishedArtifact(destination, *expected_identity) + published.append(publication) + try: + _link_and_claim(staged, destination, publication) + except FileExistsError as error: + raise ValueError( + f"refusing to overwrite existing artifact: {destination}" + ) from error + except OSError as error: + if error.errno == errno.EXDEV: + raise ValueError( + "staged and final artifacts must use the same filesystem" + ) from error + raise + if _path_identity(destination) != expected_identity: + publication.owned = False + raise ValueError(f"published artifact ownership changed: {destination}") + if not _quarantine_owned_path(staged, expected_identity, destination.parent): + raise ValueError(f"staged artifact ownership changed: {staged}") + if _path_identity(destination) != expected_identity: + publication.owned = False + raise ValueError(f"published artifact ownership changed: {destination}") + + +def _rollback_publications( + published: Sequence[_PublishedArtifact], +) -> list[tuple[Path, BaseException]]: + failures: list[tuple[Path, BaseException]] = [] + for publication in reversed(published): + if not publication.owned: + continue + try: + _quarantine_owned_path( + publication.destination, + (publication.expected_device, publication.expected_inode), + publication.destination.parent, + ) + except BaseException as error: + failures.append((publication.destination, error)) + return failures + + +def _annotate_rollback_failures( + original_error: BaseException, + failures: Sequence[tuple[Path, BaseException]], +) -> None: + for destination, cleanup_error in failures: + original_error.add_note( + f"rollback cleanup failed for {destination}: " + f"{type(cleanup_error).__name__}: {cleanup_error}" + ) + + +def finalize_mtp_export( + staging_root: Path, + output_path: Path, + receipt_path: Path, + staged_pte: Path, + staged_ptds: Sequence[Path], + source_receipt_path: Path | None, + evidence: Mapping[str, object], +) -> Path: + if staged_pte.name != output_path.name: + raise ValueError("Gemma4 MTP staged and final PTE names must match") + if len(staged_ptds) != 3: + raise ValueError("Gemma4 MTP export requires exactly three staged PTDs") + staged_artifacts = [staged_pte, *staged_ptds] + artifact_names = [path.name for path in staged_artifacts] + _require_unique_mtp_artifact_paths([{"path": name} for name in artifact_names]) + for path in staged_artifacts: + if path.is_symlink() or not path.is_file(): + raise ValueError(f"Gemma4 MTP staged artifact is not regular: {path}") + + role_paths: dict[str, Path] = {"pte": staged_pte} + staged_source: Path | None = None + if source_receipt_path is not None: + if source_receipt_path.is_symlink() or not source_receipt_path.is_file(): + raise ValueError( + "Gemma4 MTP source receipt must be a regular non-symlink file" + ) + _require_unique_mtp_artifact_paths( + [ + *({"path": name} for name in artifact_names), + {"path": source_receipt_path.name}, + ] + ) + staged_source = staging_root / source_receipt_path.name + if staged_source.exists() or staged_source.is_symlink(): + raise ValueError( + "Gemma4 MTP duplicate normalized artifact path: " + f"{source_receipt_path.name}" + ) + shutil.copyfile(source_receipt_path, staged_source) + role_paths["source"] = staged_source + + receipt = create_mtp_manifest(staging_root, role_paths, staged_ptds) + receipt["evidence"] = copy.deepcopy(evidence) + validate_mtp_manifest(staging_root, receipt) + + artifact_root = output_path.parent + artifact_root.mkdir(parents=True, exist_ok=True) + receipt_path.parent.mkdir(parents=True, exist_ok=True) + + publications = [(path, artifact_root / path.name) for path in staged_ptds] + if staged_source is not None: + publications.append((staged_source, artifact_root / staged_source.name)) + publications.append((staged_pte, output_path)) + + with tempfile.TemporaryDirectory( + prefix=f".{receipt_path.stem}.", dir=receipt_path.parent + ) as receipt_staging_directory: + staged_receipt = Path(receipt_staging_directory) / receipt_path.name + staged_receipt.write_text( + json.dumps(receipt, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + published: list[_PublishedArtifact] = [] + try: + for staged, destination in publications: + _publish_no_clobber(staged, destination, published) + validate_mtp_manifest(artifact_root, receipt) + _publish_no_clobber(staged_receipt, receipt_path, published) + validate_mtp_manifest(artifact_root, _load_json(receipt_path)) + except BaseException as error: + _annotate_rollback_failures(error, _rollback_publications(published)) + raise + return receipt_path + + def _role_paths(values: Sequence[str]) -> dict[str, Path]: result: dict[str, Path] = {} for value in values: @@ -898,6 +2912,9 @@ def _handle_closure_creation(args: argparse.Namespace) -> bool: if args.command == "create-wgsl-manifest": _write_json(args.output, create_wgsl_manifest(args.backend_root)) return True + if args.command == "create-build-recipe": + _write_json(args.output, canonical_build_recipe(args.model, args.flavor)) + return True if args.command == "create-source-receipt": _write_json( args.output, @@ -906,7 +2923,51 @@ def _handle_closure_creation(args: argparse.Namespace) -> bool: ), ) return True - return False + if args.command != "create-runtime-source": + return False + receipt = create_runtime_source_receipt( + fbsource_root=args.fbsource_root, + oss_root=args.oss_root, + backend_root=args.backend_root, + source_manifest_path=args.source_manifest, + wgsl_manifest_path=args.wgsl_manifest, + manifest_paths={"mtp": args.mtp_manifest, "plain": args.plain_manifest}, + model_roots={"mtp": args.mtp_root, "plain": args.plain_root}, + runtime_paths={ + "mtp": { + "profile": { + "javascript": args.mtp_profile_javascript, + "wasm": args.mtp_profile_wasm, + }, + "wall": { + "javascript": args.mtp_wall_javascript, + "wasm": args.mtp_wall_wasm, + }, + }, + "plain": { + "profile": { + "javascript": args.plain_profile_javascript, + "wasm": args.plain_profile_wasm, + }, + "wall": { + "javascript": args.plain_wall_javascript, + "wasm": args.plain_wall_wasm, + }, + }, + }, + build_command_paths={ + "mtp": { + "profile": args.mtp_profile_recipe, + "wall": args.mtp_wall_recipe, + }, + "plain": { + "profile": args.plain_profile_recipe, + "wall": args.plain_wall_recipe, + }, + }, + ) + _write_json(args.output, receipt) + return True def main(argv: Sequence[str] | None = None) -> int: @@ -914,17 +2975,22 @@ def main(argv: Sequence[str] | None = None) -> int: subparsers = parser.add_subparsers(dest="command", required=True) acquisition = subparsers.add_parser("validate-acquisition") acquisition.add_argument("--checkpoint-root", type=Path, required=True) + assistant_acquisition = subparsers.add_parser("validate-assistant-acquisition") + assistant_acquisition.add_argument("--checkpoint-root", type=Path, required=True) create = subparsers.add_parser("create") create.add_argument("--root", type=Path, required=True) create.add_argument("--output", type=Path, required=True) create.add_argument("--role", action="append", default=[]) create.add_argument("--ptd", action="append", type=Path, default=[]) + create_mtp = subparsers.add_parser("create-mtp") + create_mtp.add_argument("--root", type=Path, required=True) + create_mtp.add_argument("--output", type=Path, required=True) + create_mtp.add_argument("--role", action="append", default=[]) + create_mtp.add_argument("--ptd", action="append", type=Path, default=[]) + create_mtp.add_argument("--evidence", type=Path, required=True) validate = subparsers.add_parser("validate") validate.add_argument("--root", type=Path, required=True) validate.add_argument("--manifest", type=Path, required=True) - validate_mtp = subparsers.add_parser("validate-mtp") - validate_mtp.add_argument("--root", type=Path, required=True) - validate_mtp.add_argument("--manifest", type=Path, required=True) create_source = subparsers.add_parser("create-source-manifest") create_source.add_argument("--fbsource-root", type=Path, required=True) create_source.add_argument("--oss-root", type=Path, required=True) @@ -937,6 +3003,63 @@ def main(argv: Sequence[str] | None = None) -> int: create_source_receipt.add_argument("--oss-root", type=Path, required=True) create_source_receipt.add_argument("--backend-root", type=Path, required=True) create_source_receipt.add_argument("--output", type=Path, required=True) + validate_mtp = subparsers.add_parser("validate-mtp") + validate_mtp.add_argument("--root", type=Path, required=True) + validate_mtp.add_argument("--manifest", type=Path, required=True) + stage_runtime = subparsers.add_parser("stage-runtime") + stage_runtime.add_argument("--destination-root", type=Path, required=True) + stage_runtime.add_argument("--plain-root", type=Path, required=True) + stage_runtime.add_argument("--plain-receipt", type=Path, required=True) + stage_runtime.add_argument("--mtp-root", type=Path, required=True) + stage_runtime.add_argument("--mtp-receipt", type=Path, required=True) + stage_runtime.add_argument("--runtime-source-receipt", type=Path, required=True) + stage_runtime.add_argument("--target-prefill-receipt", type=Path, required=True) + stage_runtime.add_argument("--plain-profile-javascript", type=Path, required=True) + stage_runtime.add_argument("--plain-profile-wasm", type=Path, required=True) + stage_runtime.add_argument("--plain-wall-javascript", type=Path, required=True) + stage_runtime.add_argument("--plain-wall-wasm", type=Path, required=True) + stage_runtime.add_argument("--mtp-wall-javascript", type=Path, required=True) + stage_runtime.add_argument("--mtp-wall-wasm", type=Path, required=True) + stage_runtime.add_argument("--mtp-profile-javascript", type=Path, required=True) + stage_runtime.add_argument("--mtp-profile-wasm", type=Path, required=True) + stage_runtime.add_argument("--source-manifest", type=Path, required=True) + stage_runtime.add_argument("--wgsl-manifest", type=Path, required=True) + stage_runtime.add_argument("--plain-profile-recipe", type=Path, required=True) + stage_runtime.add_argument("--plain-wall-recipe", type=Path, required=True) + stage_runtime.add_argument("--mtp-wall-recipe", type=Path, required=True) + stage_runtime.add_argument("--mtp-profile-recipe", type=Path, required=True) + create_runtime = subparsers.add_parser("create-runtime-source") + create_runtime.add_argument("--output", type=Path, required=True) + create_runtime.add_argument("--fbsource-root", type=Path, required=True) + create_runtime.add_argument("--oss-root", type=Path, required=True) + create_runtime.add_argument("--backend-root", type=Path, required=True) + create_runtime.add_argument("--source-manifest", type=Path, required=True) + create_runtime.add_argument("--wgsl-manifest", type=Path, required=True) + create_runtime.add_argument("--plain-manifest", type=Path, required=True) + create_runtime.add_argument("--mtp-manifest", type=Path, required=True) + create_runtime.add_argument("--plain-root", type=Path, required=True) + create_runtime.add_argument("--mtp-root", type=Path, required=True) + create_runtime.add_argument("--plain-profile-javascript", type=Path, required=True) + create_runtime.add_argument("--plain-profile-wasm", type=Path, required=True) + create_runtime.add_argument("--plain-profile-recipe", type=Path, required=True) + create_runtime.add_argument("--plain-wall-javascript", type=Path, required=True) + create_runtime.add_argument("--plain-wall-wasm", type=Path, required=True) + create_runtime.add_argument("--plain-wall-recipe", type=Path, required=True) + create_runtime.add_argument("--mtp-wall-javascript", type=Path, required=True) + create_runtime.add_argument("--mtp-wall-wasm", type=Path, required=True) + create_runtime.add_argument("--mtp-wall-recipe", type=Path, required=True) + create_runtime.add_argument("--mtp-profile-javascript", type=Path, required=True) + create_runtime.add_argument("--mtp-profile-wasm", type=Path, required=True) + create_runtime.add_argument("--mtp-profile-recipe", type=Path, required=True) + create_recipe = subparsers.add_parser("create-build-recipe") + create_recipe.add_argument( + "--model", choices=sorted(RUNTIME_BUILD_TARGETS), required=True + ) + create_recipe.add_argument("--flavor", choices=("profile", "wall"), required=True) + create_recipe.add_argument("--output", type=Path, required=True) + validate_runtime = subparsers.add_parser("validate-runtime") + validate_runtime.add_argument("--root", type=Path, required=True) + validate_runtime.add_argument("--manifest", type=Path, required=True) args = parser.parse_args(argv) if args.command == "validate-acquisition": @@ -944,10 +3067,48 @@ def main(argv: Sequence[str] | None = None) -> int: return 0 if _handle_closure_creation(args): return 0 - if args.command == "create": - manifest = create_plain_manifest( - args.root, _role_paths(args.role), args.ptd + if args.command == "validate-assistant-acquisition": + validate_assistant_export_identity(args.checkpoint_root) + return 0 + if args.command == "stage-runtime": + stage_combined_runtime( + args.destination_root, + args.plain_root, + args.plain_receipt, + args.mtp_root, + args.mtp_receipt, + args.runtime_source_receipt, + args.mtp_wall_javascript, + args.mtp_wall_wasm, + args.mtp_profile_javascript, + args.mtp_profile_wasm, + plain_profile_javascript_path=args.plain_profile_javascript, + plain_profile_wasm_path=args.plain_profile_wasm, + plain_wall_javascript_path=args.plain_wall_javascript, + plain_wall_wasm_path=args.plain_wall_wasm, + source_manifest_path=args.source_manifest, + wgsl_manifest_path=args.wgsl_manifest, + build_recipe_paths={ + "mtp": { + "profile": args.mtp_profile_recipe, + "wall": args.mtp_wall_recipe, + }, + "plain": { + "profile": args.plain_profile_recipe, + "wall": args.plain_wall_recipe, + }, + }, + target_prefill_receipt_path=args.target_prefill_receipt, ) + return 0 + if args.command in {"create", "create-mtp"}: + create_fn = ( + create_plain_manifest if args.command == "create" else create_mtp_manifest + ) + manifest = create_fn(args.root, _role_paths(args.role), args.ptd) + if args.command == "create-mtp": + manifest["evidence"] = _load_json(args.evidence) + validate_mtp_manifest(args.root, manifest) args.output.write_text( json.dumps(manifest, indent=2, sort_keys=True) + "\n", encoding="utf-8", @@ -955,7 +3116,12 @@ def main(argv: Sequence[str] | None = None) -> int: return 0 manifest = _load_json(args.manifest) - validate_plain_manifest(args.root, manifest) + if args.command == "validate": + validate_plain_manifest(args.root, manifest) + elif args.command == "validate-mtp": + validate_mtp_manifest(args.root, manifest) + else: + validate_combined_runtime_envelope(args.root, manifest) return 0