[ExecuTorch][WebGPU] Add Gemma 4 speculative decode runtime

backends/webgpu/CMakeLists.txt

@@ -1,3 +1,3 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
 #
@@ -117,6 +117,19 @@
 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 @@
     gemma4_plain_wasm PRIVATE $<BUILD_INTERFACE:${EXECUTORCH_ROOT}/..>
   )
   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 @@
                "${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 $<BUILD_INTERFACE:${EXECUTORCH_ROOT}/..>
+  )
+  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(

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()

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

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()