diff --git a/backends/webgpu/CMakeLists.txt b/backends/webgpu/CMakeLists.txt index adeea976fac..0298fdd5a3c 100644 --- a/backends/webgpu/CMakeLists.txt +++ b/backends/webgpu/CMakeLists.txt @@ -116,6 +116,51 @@ executorch_target_link_options_shared_lib(webgpu_backend) set_property(TARGET webgpu_backend PROPERTY CXX_STANDARD 17) +if(EMSCRIPTEN) + add_executable( + gemma4_plain_wasm + ${EXECUTORCH_ROOT}/examples/models/gemma4/runner/gemma4_plain_wasm.cpp + ) + target_include_directories( + gemma4_plain_wasm PRIVATE $ + ) + target_link_libraries( + gemma4_plain_wasm PRIVATE webgpu_backend webgpu_model_loader extension_tensor + ) + target_compile_options( + gemma4_plain_wasm PRIVATE -fexceptions "--use-port=emdawnwebgpu" + ) + if(EXECUTORCH_BUILD_WEBGPU_PROFILING) + target_compile_definitions( + gemma4_plain_wasm PRIVATE WGPU_BACKEND_ENABLE_PROFILING + ) + endif() + target_link_options( + gemma4_plain_wasm + 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_profile_enable,_et_profile,_et_get_last_prefill_token_count,_et_get_route_contract_version,_et_get_last_route_mask,_et_get_last_route_conflict_count,_malloc,_free" + "-sEXPORTED_RUNTIME_METHODS=ccall,cwrap,FS,HEAP32" + -sSTACK_SIZE=8388608 + -sASYNCIFY_STACK_SIZE=1048576 + -sMODULARIZE=1 + -sEXPORT_NAME=createWebGPULlama + ) + set_target_properties( + gemma4_plain_wasm + PROPERTIES OUTPUT_NAME "webgpu_llama" + RUNTIME_OUTPUT_DIRECTORY + "${CMAKE_CURRENT_BINARY_DIR}/browser_gemma4_plain" + CXX_STANDARD 17 + ) +endif() + install( TARGETS webgpu_backend EXPORT ExecuTorchTargets diff --git a/examples/models/gemma4/BUCK b/examples/models/gemma4/BUCK index 19f0ff90c93..9b57f93bc56 100644 --- a/examples/models/gemma4/BUCK +++ b/examples/models/gemma4/BUCK @@ -1,6 +1,6 @@ load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target", "non_fbcode_target") load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") -load(":targets.bzl", "define_common_targets") +load(":targets.bzl", "define_common_targets", "define_webgpu_python_targets") oncall("executorch") @@ -8,6 +8,8 @@ non_fbcode_target(_kind = define_common_targets,) fbcode_target(_kind = define_common_targets,) +define_webgpu_python_targets() + # Text decoder module fbcode_target(_kind = runtime.python_library, name = "text_decoder", @@ -77,6 +79,7 @@ fbcode_target(_kind = runtime.python_binary, "//executorch/kernels/quantized:aot_lib", ], deps = [ + ":webgpu_support", ":text_decoder", ":speech_transform", ":quant_utils", diff --git a/examples/models/gemma4/README.md b/examples/models/gemma4/README.md index 686c60db8ac..7289978f0c3 100644 --- a/examples/models/gemma4/README.md +++ b/examples/models/gemma4/README.md @@ -45,6 +45,84 @@ buck2 run fbcode//executorch/examples/models/gemma4:export_gemma4 -- \ --checkpoint_path /tmp/gemma4-e2b-it --no-audio ``` +### Plain E2B WebGPU + +The WebGPU path is independently exportable and text-only. It preserves an +8960-token KV capacity while bounding each input call to 512 tokens, returns a +delegated `Long[1, 1]` greedy token, and splits external constants into three +ordered PTD files below the browser binding/fetch limit. The default XNNPACK +export and runner are unchanged. + +Acquire the checkpoint from +`google/gemma-4-E2B-it-qat-q4_0-unquantized` at immutable revision +`6befbaca7398925921802abd1f277b495b78b738`, then validate every staged byte: + +```bash +hf download google/gemma-4-E2B-it-qat-q4_0-unquantized \ + model.safetensors config.json tokenizer.json tokenizer_config.json \ + generation_config.json processor_config.json chat_template.jinja \ + README.md .gitattributes \ + --revision 6befbaca7398925921802abd1f277b495b78b738 \ + --local-dir /tmp/gemma4-e2b-it +``` + +```bash +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: + +```bash +: "${FBSOURCE_ROOT:?set the clean fbsource checkout root}" +: "${OSS_ROOT:?set the clean ExecuTorch OSS checkout root}" +python -m executorch.examples.models.gemma4.webgpu_artifact_manifest \ + create-source-manifest --fbsource-root "$FBSOURCE_ROOT" \ + --oss-root "$OSS_ROOT" \ + --output /tmp/gemma4-source-manifest.json +python -m executorch.examples.models.gemma4.webgpu_artifact_manifest \ + create-wgsl-manifest \ + --backend-root "$FBSOURCE_ROOT/xplat/executorch/backends/webgpu" \ + --output /tmp/gemma4-wgsl-manifest.json +python -m executorch.examples.models.gemma4.webgpu_artifact_manifest \ + create-source-receipt \ + --fbsource-root "$FBSOURCE_ROOT" --oss-root "$OSS_ROOT" \ + --backend-root "$FBSOURCE_ROOT/xplat/executorch/backends/webgpu" \ + --output /tmp/gemma4-source-receipt.json +``` + +Export plain Gemma 4: + +```bash +buck2 run fbcode//executorch/examples/models/gemma4:export_gemma4 -- \ + --checkpoint_path /tmp/gemma4-e2b-it \ + --output_path /tmp/gemma4-webgpu/model.pte \ + --backend webgpu --quantize 8da4w+emb4 \ + --max_seq_len 8960 --max_input_len 512 \ + --no-audio --no-vision \ + --source_receipt_path /tmp/gemma4-source-receipt.json \ + --artifact_manifest_output /tmp/gemma4-e2b-webgpu.json +``` + +The source-closure gate writes `gemma4-source-receipt.json` before this export. +The exporter reads the actual tensor-data insertion order, writes all three +content-named PTDs, and creates the manifest without renaming or globbing them. +Keep the manifest output outside the flat artifact staging directory. + +```bash +buck2 run fbcode//executorch/examples/models/gemma4:webgpu_artifact_manifest -- \ + validate --root /tmp/gemma4-webgpu \ + --manifest /tmp/gemma4-e2b-webgpu.json +``` + +`manifests/gemma4_e2b_webgpu.json` pins the accepted ctx8960 behavior-oracle +quartet, but marks its old-worktree source closure as pending. It cannot satisfy +the production validator because it has no final-source receipt. Rebuild it +from the reviewed stack before claiming source-current performance or +reproduction. Dashboard and internal publication paths are evidence only and +are never source dependencies. + ## Model Variants Default export includes all modalities (audio + vision + text). Default context length: 1024 tokens (`--max_seq_len`). diff --git a/examples/models/gemma4/export_gemma4.py b/examples/models/gemma4/export_gemma4.py index d59d6c82615..077c58c298c 100644 --- a/examples/models/gemma4/export_gemma4.py +++ b/examples/models/gemma4/export_gemma4.py @@ -27,6 +27,7 @@ import argparse import functools import gc +import json import logging from pathlib import Path @@ -36,6 +37,16 @@ logger = logging.getLogger(__name__) +class _GreedyTokenOutput(torch.nn.Module): + def __init__(self, model: torch.nn.Module) -> None: + super().__init__() + self.model = model + + def forward(self, input_ids: torch.Tensor, input_pos: torch.Tensor) -> torch.Tensor: + logits = self.model(input_ids, input_pos=input_pos) + return torch.argmax(logits, dim=-1).to(torch.long) + + def _export_speech_transform(checkpoint_path: str): """Export speech transform model. Returns ExportedProgram.""" from executorch.examples.models.gemma4.speech_transform import ( @@ -452,6 +463,9 @@ def _export_text_decoder( variant: str = "e2b", quantize_kv_cache: bool = False, use_custom_sdpa: bool = True, + backend: str = "xnnpack", + max_input_len: int = 512, + compact_output: bool = True, ): """Export text decoder. Returns ExportedProgram.""" from executorch.examples.models.gemma4.quant_utils import ( @@ -523,19 +537,29 @@ def _export_text_decoder( model = apply_linear_quantization(model, linear_quant, group_size=group_size) model.eval() - # Export with audio embeds + dynamic shapes - example_inputs = model_wrapper.get_example_inputs_with_audio(seq_len=770) - dynamic_shapes = model_wrapper.get_dynamic_shapes(with_audio_embeds=True) + if backend == "webgpu": + example_inputs = model_wrapper.get_example_inputs_prefill(seq_len=max_input_len) + dynamic_shapes = model_wrapper.get_dynamic_shapes(max_input_len=max_input_len) + if compact_output: + model = _GreedyTokenOutput(model) + model.eval() + args = (example_inputs[0],) + kwargs = {"input_pos": example_inputs[1]} + else: + example_inputs = model_wrapper.get_example_inputs_with_audio(seq_len=770) + dynamic_shapes = model_wrapper.get_dynamic_shapes(with_audio_embeds=True) + args = (example_inputs[0],) + kwargs = { + "input_pos": example_inputs[1], + "inputs_embeds": example_inputs[2], + } with torch.nn.attention.sdpa_kernel([torch.nn.attention.SDPBackend.MATH]): with torch.no_grad(): ep = torch.export.export( model, - (example_inputs[0],), - kwargs={ - "input_pos": example_inputs[1], - "inputs_embeds": example_inputs[2], - }, + args, + kwargs=kwargs, dynamic_shapes=dynamic_shapes, ) @@ -559,6 +583,9 @@ def _export_components( include_audio: bool, include_vision: bool, use_custom_sdpa: bool, + backend: str = "xnnpack", + max_input_len: int = 512, + compact_output: bool = True, ) -> dict: """Export each requested component to an ExportedProgram.""" components = [] @@ -605,19 +632,34 @@ def _export_components( variant=variant, quantize_kv_cache=quantize_kv_cache, use_custom_sdpa=use_custom_sdpa, + backend=backend, + max_input_len=max_input_len, + compact_output=compact_output, ) return programs -def _build_partitioners( +def build_partitioners( + backend: str, include_audio: bool, include_vision: bool, audio_quantize: str, vision_quantize: str, text_quantize: str, ) -> dict: - """Build per-method XNNPACK partitioner lists.""" + """Build per-method partitioners for the selected backend.""" + if backend == "webgpu": + from executorch.examples.models.gemma4.webgpu_partitioner import ( + build_webgpu_partitioner, + ) + + return { + "text_decoder": [build_webgpu_partitioner(text_quantize)], + } + if backend != "xnnpack": + raise ValueError(f"Unsupported backend: {backend}") + from executorch.backends.xnnpack.partition.xnnpack_partitioner import ( XnnpackDynamicallyQuantizedPartitioner, XnnpackPartitioner, @@ -639,8 +681,14 @@ def _for(quantize: str) -> list: return partitioners -def _build_transform_passes(include_audio: bool, include_vision: bool) -> dict: - """Build per-method transform passes (text decoder gets bitwise lowering).""" +def _build_transform_passes( + include_audio: bool, include_vision: bool, backend: str = "xnnpack" +) -> dict: + """Build per-method transform passes (text decoder gets bitwise lowering). + + WebGPU additionally runs its edge rewrites here: the partitioner may not + mutate the graph, so they must land before partitioning. + """ from executorch.exir.dialects._ops import ops as exir_ops from executorch.exir.pass_base import ExportPass @@ -681,9 +729,68 @@ def call_operator(self, op, args, kwargs, meta): if include_vision: transform_passes["vision_encoder"] = [] transform_passes["text_decoder"] = [_ReplaceBitwiseScalarPass(bitwise_ops)] + if backend == "webgpu": + from executorch.examples.models.gemma4.webgpu_partitioner import ( + build_webgpu_transform_passes, + ) + + transform_passes["text_decoder"].extend(build_webgpu_transform_passes()) return transform_passes +def _rewrite_webgpu_negative_select_as_symint( + program: torch.export.ExportedProgram, +) -> int: + """Preserve a negative input select across decomposition.""" + import executorch.backends.vulkan.custom_ops_lib # noqa: F401 + + graph = program.graph_module.graph + item_targets = { + torch.ops.aten.item.default, + torch.ops.aten._local_scalar_dense.default, + } + replacements = 0 + for item_node in list(graph.nodes): + if item_node.op != "call_function" or item_node.target not in item_targets: + continue + if len(item_node.args) != 1 or not isinstance(item_node.args[0], torch.fx.Node): + continue + select_node = item_node.args[0] + if ( + select_node.op != "call_function" + or select_node.target != torch.ops.aten.select.int + or len(select_node.args) < 3 + or not isinstance(select_node.args[2], int) + or select_node.args[2] >= 0 + ): + continue + source = select_node.args[0] + if not isinstance(source, torch.fx.Node) or source.op != "placeholder": + raise ValueError( + "WebGPU negative select_as_symint source must be a graph input" + ) + source_dtype = getattr(source.meta.get("val"), "dtype", None) + if source_dtype not in {torch.int32, torch.int64}: + raise ValueError( + "WebGPU negative select_as_symint requires an integral input" + ) + with graph.inserting_before(item_node): + replacement = graph.call_function( + torch.ops.et_vk.select_as_symint.default, + args=(source, select_node.args[1], select_node.args[2]), + ) + replacement.meta = item_node.meta.copy() + item_node.replace_all_uses_with(replacement) + graph.erase_node(item_node) + if not select_node.users: + graph.erase_node(select_node) + replacements += 1 + + if replacements: + program.graph_module.recompile() + return replacements + + def export_single_pte( checkpoint_path: str, output_path: str, @@ -699,6 +806,11 @@ def export_single_pte( include_audio: bool = True, include_vision: bool = True, use_custom_sdpa: bool = True, + backend: str = "xnnpack", + max_input_len: int = 512, + compact_output: bool = True, + artifact_manifest_output: str | None = None, + source_receipt_path: str | None = None, ) -> Path: """Export components into a single PTE. @@ -713,6 +825,37 @@ def export_single_pte( ConstraintBasedSymShapeEvalPass, ) + if backend == "webgpu": + if variant != "e2b": + raise ValueError("WebGPU export currently supports only Gemma 4 E2B") + if include_audio or include_vision: + raise ValueError("WebGPU export is text-only; disable audio and vision") + if not use_custom_sdpa: + raise ValueError("WebGPU export requires Gemma 4 custom SDPA") + if max_seq_len != 8960 or max_input_len != 512: + raise ValueError( + "WebGPU production export requires max_seq_len=8960 and " + "max_input_len=512" + ) + if text_quantize != "8da4w+emb4" or group_size != 128: + raise ValueError( + "WebGPU production export requires 8da4w+emb4 with group_size=128" + ) + if tied_embedding or quantize_kv_cache or not compact_output: + raise ValueError( + "WebGPU production export requires untied embeddings, fp32 KV cache, " + "and compact greedy output" + ) + from executorch.examples.models.gemma4.webgpu_artifact_manifest import ( + validate_export_identity, + ) + + validate_export_identity(Path(checkpoint_path)) + if (artifact_manifest_output is None) != (source_receipt_path is None): + raise ValueError( + "WebGPU artifact manifest output and source receipt must be provided together" + ) + programs = _export_components( checkpoint_path=checkpoint_path, text_quantize=text_quantize, @@ -727,16 +870,34 @@ def export_single_pte( include_audio=include_audio, include_vision=include_vision, use_custom_sdpa=use_custom_sdpa, + backend=backend, + max_input_len=max_input_len, + compact_output=compact_output, ) + if backend == "webgpu": + for name, program in programs.items(): + replacements = _rewrite_webgpu_negative_select_as_symint(program) + if replacements: + logger.info( + "Rewrote %d negative select-as-SymInt path(s) in %s", + replacements, + name, + ) + logger.info("Combining into single PTE...") for name in programs: programs[name] = programs[name].run_decompositions({}) - partitioners = _build_partitioners( - include_audio, include_vision, audio_quantize, vision_quantize, text_quantize + partitioners = build_partitioners( + backend, + include_audio, + include_vision, + audio_quantize, + vision_quantize, + text_quantize, ) - transform_passes = _build_transform_passes(include_audio, include_vision) + transform_passes = _build_transform_passes(include_audio, include_vision, backend) edge_manager = to_edge_transform_and_lower( programs, @@ -753,7 +914,7 @@ def export_single_pte( ) ) - output_path = Path(output_path) + output_path = Path(output_path).resolve() output_path.parent.mkdir(parents=True, exist_ok=True) with open(output_path, "wb") as f: et_program.write_to_file(f) @@ -763,6 +924,45 @@ def export_single_pte( et_program.write_tensor_data_to_file(tensor_data_dir) logger.info(f"Tensor data written to: {tensor_data_dir}") + if backend == "webgpu" and artifact_manifest_output is not None: + assert source_receipt_path is not None + from executorch.examples.models.gemma4.webgpu_artifact_manifest import ( + create_plain_manifest, + validate_plain_manifest, + ) + + manifest_output = Path(artifact_manifest_output).resolve() + try: + manifest_output.relative_to(output_path.parent) + except ValueError: + pass + else: + raise ValueError( + "Artifact manifest output must be outside the staging root" + ) + tensor_data = et_program._tensor_data + if tensor_data is None: + raise ValueError("WebGPU export did not produce external constants") + ptd_paths = [ + Path(name if name.endswith(".ptd") else f"{name}.ptd") + for name in tensor_data + ] + manifest = create_plain_manifest( + output_path.parent, + { + "pte": output_path, + "source": Path(source_receipt_path).resolve(), + }, + ptd_paths, + ) + validate_plain_manifest(output_path.parent, manifest) + manifest_output.parent.mkdir(parents=True, exist_ok=True) + manifest_output.write_text( + json.dumps(manifest, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + logger.info(f"WebGPU artifact manifest written to: {manifest_output}") + size_mb = output_path.stat().st_size / 1024 / 1024 logger.info(f"Single PTE exported: {output_path} ({size_mb:.1f} MB)") return output_path @@ -783,6 +983,13 @@ def main(): choices=["e2b", "e4b"], help="Model variant (default: e2b)", ) + parser.add_argument( + "--backend", + type=str, + default="xnnpack", + choices=["xnnpack", "webgpu"], + help="Delegate backend (default: xnnpack)", + ) parser.add_argument( "--output_path", type=str, @@ -795,6 +1002,12 @@ def main(): default=1024, help="Maximum sequence length for text decoder KV cache", ) + parser.add_argument( + "--max_input_len", + type=int, + default=512, + help="Maximum per-call input length for dynamic WebGPU export", + ) parser.add_argument( "--quantize", type=str, @@ -860,6 +1073,22 @@ def main(): help="Route attention through llama::custom_sdpa (tiled flash attention). " "Pass --no-use_custom_sdpa to fall back to matmul attention.", ) + parser.add_argument( + "--compact_output", + action=argparse.BooleanOptionalAction, + default=True, + help="Return the Long[1,1] greedy token for WebGPU export.", + ) + parser.add_argument( + "--artifact_manifest_output", + type=str, + help="Write a byte-validated WebGPU manifest outside the artifact root.", + ) + parser.add_argument( + "--source_receipt_path", + type=str, + help="Final-source receipt staged beside the WebGPU PTE and PTDs.", + ) args = parser.parse_args() export_single_pte( @@ -877,6 +1106,11 @@ def main(): include_audio=not args.no_audio, include_vision=not args.no_vision, use_custom_sdpa=args.use_custom_sdpa, + backend=args.backend, + max_input_len=args.max_input_len, + compact_output=args.compact_output, + artifact_manifest_output=args.artifact_manifest_output, + source_receipt_path=args.source_receipt_path, ) diff --git a/examples/models/gemma4/manifests/gemma4_e2b_webgpu.json b/examples/models/gemma4/manifests/gemma4_e2b_webgpu.json new file mode 100644 index 00000000000..581cc7b4cbd --- /dev/null +++ b/examples/models/gemma4/manifests/gemma4_e2b_webgpu.json @@ -0,0 +1,124 @@ +{ + "acquisition": { + "files": { + ".gitattributes": { + "bytes": 1570, + "sha256": "34448b82c17d60fec9b65b1f093c115ddbaadc04beb1b0140b6bfed2e012a930" + }, + "README.md": { + "bytes": 29351, + "sha256": "aaab87052837925e0fb400bb20700553b11088fa5b3ae21fa0c1ec5da53637a4" + }, + "chat_template.jinja": { + "bytes": 18569, + "sha256": "0a2c8073c878ab1da004bee933a998606537bbb62016310352c7285c3f01c5b5" + }, + "config.json": { + "bytes": 4946, + "sha256": "bbeff1e2fd3fe282536e7ace02309d43e0dbd9b6ac4b6a149b97e3ab6942a878" + }, + "generation_config.json": { + "bytes": 203, + "sha256": "b69207f9be617e982d13cc273cce6fd88c98dda99a4bdc5e2d52ffe0a0d9f0a9" + }, + "model.safetensors": { + "bytes": 10208852878, + "sha256": "33fe0cece08fb527ffefbd1a3a9ce73bd71073727993a283506293e5c6bf0137" + }, + "processor_config.json": { + "bytes": 1689, + "sha256": "32bdf45d2ad4cc29a0822ddd157a182de76644f0419a6228d151495256e9813c" + }, + "tokenizer.json": { + "bytes": 32169626, + "sha256": "cc8d3a0ce36466ccc1278bf987df5f71db1719b9ca6b4118264f45cb627bfe0f" + }, + "tokenizer_config.json": { + "bytes": 3729, + "sha256": "3ab5c7b94dc97d65ca7064496fa69b88ff875378e1cb7ee3e43070c3a8170999" + } + }, + "repo_id": "google/gemma-4-E2B-it-qat-q4_0-unquantized", + "revision": "6befbaca7398925921802abd1f277b495b78b738" + }, + "artifacts": [ + { + "bytes": 697820048, + "path": "model.pte", + "role": "pte", + "sha256": "3c8cd7cffe28b7a5a7f2a822d61cb08ca75e2094296c365026de8c0673b9bb33" + }, + { + "bytes": 490343936, + "path": "vulkan_constants_5744ffbc74f9b1f05c9b1de0242271e80917b06d119a752ab06357a3ce34bac0.ptd", + "role": "ptd", + "sha256": "46bde3894dde6a31c6988158f764bb819f1355b17a88d76c916d0ae632656980" + }, + { + "bytes": 1499632256, + "path": "vulkan_constants_c99ae26962b09639368ed489fd7383e97c2c1b2b65cb0c49957f7657b38f4235.ptd", + "role": "ptd", + "sha256": "82a18ddf74f7befdeab3a946cfb2e0ff4d79f52659ca9832b59dc37a3c45140e" + }, + { + "bytes": 1396131456, + "path": "vulkan_constants_dc47af467f37efd71f907f227e6414712f1e2c3419daa2692b29ed92687110b7.ptd", + "role": "ptd", + "sha256": "8dfc0f10781f98fd62cf7bc59eeb071f435df2403fb7ebcb169e1df7669217d2" + } + ], + "export": { + "backend": "webgpu", + "max_input_len": 512, + "max_seq_len": 8960, + "methods": [ + "text_decoder" + ], + "output": { + "dtype": "int64", + "semantic": "greedy_token", + "shape": [ + 1, + 1 + ] + }, + "quantization": "8da4w+emb4" + }, + "model": { + "architecture": { + "global_head_dim": 512, + "head_dim": 256, + "hidden_size": 1536, + "intermediate_size": 6144, + "layer_types": [ + "sliding_attention", "sliding_attention", "sliding_attention", "sliding_attention", "full_attention", + "sliding_attention", "sliding_attention", "sliding_attention", "sliding_attention", "full_attention", + "sliding_attention", "sliding_attention", "sliding_attention", "sliding_attention", "full_attention", + "sliding_attention", "sliding_attention", "sliding_attention", "sliding_attention", "full_attention", + "sliding_attention", "sliding_attention", "sliding_attention", "sliding_attention", "full_attention", + "sliding_attention", "sliding_attention", "sliding_attention", "sliding_attention", "full_attention", + "sliding_attention", "sliding_attention", "sliding_attention", "sliding_attention", "full_attention" + ], + "model_type": "gemma4", + "num_attention_heads": 8, + "num_hidden_layers": 35, + "num_key_value_heads": 1, + "num_kv_shared_layers": 20, + "vocab_size": 262144 + }, + "source_config": { + "path": "config/e2b_config.json", + "sha256": "526e9fd34a8a489c35952535335a4b8556e9169d851187a895f80286e7466206" + } + }, + "provenance": { + "artifact_status": "accepted_behavior_oracle", + "source_closure": "pending_final_source_rebuild" + }, + "ptd_order": [ + "vulkan_constants_5744ffbc74f9b1f05c9b1de0242271e80917b06d119a752ab06357a3ce34bac0.ptd", + "vulkan_constants_c99ae26962b09639368ed489fd7383e97c2c1b2b65cb0c49957f7657b38f4235.ptd", + "vulkan_constants_dc47af467f37efd71f907f227e6414712f1e2c3419daa2692b29ed92687110b7.ptd" + ], + "schema_version": 1 +} diff --git a/examples/models/gemma4/runner/gemma4_plain_wasm.cpp b/examples/models/gemma4/runner/gemma4_plain_wasm.cpp new file mode 100644 index 00000000000..0b040bb4058 --- /dev/null +++ b/examples/models/gemma4/runner/gemma4_plain_wasm.cpp @@ -0,0 +1,394 @@ +/* + * 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 +#ifdef WGPU_BACKEND_ENABLE_PROFILING +#include +#endif +#include + +#ifdef __EMSCRIPTEN__ +#include +#define GEMMA4_WASM_EXPORT EMSCRIPTEN_KEEPALIVE +#else +#define GEMMA4_WASM_EXPORT +#endif + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace { + +using executorch::aten::ScalarType; +using executorch::backends::webgpu::WebGPUContext; +using executorch::backends::webgpu::WebGPUModelLoadSpec; +using executorch::backends::webgpu::compare_and_set_default_webgpu_context; +using executorch::backends::webgpu::create_webgpu_context; +using executorch::backends::webgpu::destroy_webgpu_context; +#ifdef WGPU_BACKEND_ENABLE_PROFILING +using executorch::backends::webgpu::g_last_route_conflict_count; +using executorch::backends::webgpu::g_last_route_mask; +#endif +using executorch::backends::webgpu::get_default_webgpu_context; +using executorch::backends::webgpu::load_webgpu_model; +using executorch::extension::Module; +using executorch::extension::make_tensor_ptr; +using executorch::runtime::Error; +using executorch::runtime::EValue; +using executorch::runtime::MethodMeta; +using executorch::runtime::Tag; +using executorch::runtime::TensorInfo; + +constexpr char kMethodName[] = "text_decoder"; +constexpr size_t kExpectedPtdCount = 3; +constexpr int kMaxInputLength = 512; +constexpr int kMaxSequenceLength = 8960; + +WebGPUContext g_context; +bool g_owns_default_context = false; +std::unique_ptr g_module; +int g_last_prefill_tokens = 0; + +void reset_runtime_observations() { + g_last_prefill_tokens = 0; +#ifdef WGPU_BACKEND_ENABLE_PROFILING + if (g_context.querypool) { + g_context.querypool->reset(0); + } + g_last_route_mask = 0; + g_last_route_conflict_count = 0; +#endif +} + +void unload_model() { + g_module.reset(); + reset_runtime_observations(); +} + +void release_context() { + if (!g_owns_default_context) { + return; + } + compare_and_set_default_webgpu_context(&g_context, nullptr); + destroy_webgpu_context(g_context); + g_owns_default_context = false; +} + +bool is_long_tensor(const executorch::runtime::Result& info) { + return info.ok() && info.get().scalar_type() == ScalarType::Long; +} + +bool validate_method_meta(const MethodMeta& meta) { + if (meta.num_inputs() != 2 || meta.num_outputs() != 1 || + !meta.uses_backend("VulkanBackend")) { + return false; + } + const auto input_ids_tag = meta.input_tag(0); + const auto input_pos_tag = meta.input_tag(1); + const auto output_tag = meta.output_tag(0); + if (!input_ids_tag.ok() || input_ids_tag.get() != Tag::Tensor || + !input_pos_tag.ok() || input_pos_tag.get() != Tag::Tensor || + !output_tag.ok() || output_tag.get() != Tag::Tensor) { + return false; + } + + const auto input_ids = meta.input_tensor_meta(0); + const auto input_pos = meta.input_tensor_meta(1); + const auto output = meta.output_tensor_meta(0); + if (!is_long_tensor(input_ids) || !is_long_tensor(input_pos) || + !is_long_tensor(output)) { + return false; + } + const auto input_ids_sizes = input_ids.get().sizes(); + const auto input_pos_sizes = input_pos.get().sizes(); + const auto output_sizes = output.get().sizes(); + return input_ids_sizes.size() == 2 && input_ids_sizes[0] == 1 && + input_pos_sizes.size() == 1 && output_sizes.size() == 2 && + output_sizes[0] == 1 && output_sizes[1] == 1 && + output.get().nbytes() == sizeof(int64_t); +} + +bool validate_loaded_method(Module& module) { + auto meta = module.method_meta(kMethodName); + return meta.ok() && validate_method_meta(meta.get()); +} + +std::vector parse_ptd_paths(const char* data_paths) { + std::vector ptd_paths; + if (data_paths == nullptr) { + return ptd_paths; + } + std::istringstream stream{std::string(data_paths)}; + std::string path; + while (std::getline(stream, path)) { + if (!path.empty()) { + ptd_paths.push_back(std::move(path)); + } + } + return ptd_paths; +} + +int read_compact_token(const std::vector& values) { + if (values.size() != 1 || !values[0].isTensor()) { + return -1; + } + const auto& output = values[0].toTensor(); + if (output.scalar_type() != ScalarType::Long || output.dim() != 2 || + output.size(0) != 1 || output.size(1) != 1 || output.numel() != 1 || + output.nbytes() != sizeof(int64_t)) { + return -1; + } + const int64_t token = output.const_data_ptr()[0]; + if (token < 0 || token > std::numeric_limits::max()) { + return -1; + } + return static_cast(token); +} + +int execute_tokens(const int* tokens, int count, int position) { + if (!g_module || tokens == nullptr || count <= 0 || + count > kMaxInputLength || position < 0 || + position > kMaxSequenceLength - count) { + return -1; + } + + std::vector input_ids(static_cast(count)); + std::vector input_pos(static_cast(count)); + for (int index = 0; index < count; ++index) { + input_ids[static_cast(index)] = tokens[index]; + input_pos[static_cast(index)] = position + index; + } + auto input_ids_tensor = make_tensor_ptr( + {1, count}, std::move(input_ids), {}, {}, ScalarType::Long); + auto input_pos_tensor = make_tensor_ptr( + {count}, std::move(input_pos), {}, {}, ScalarType::Long); + auto result = g_module->execute( + kMethodName, + {EValue(std::move(input_ids_tensor)), EValue(std::move(input_pos_tensor))}); + if (!result.ok()) { + std::printf("Gemma4 text_decoder failed: %d\n", (int)result.error()); + return -1; + } + return read_compact_token(result.get()); +} + +} // namespace + +extern "C" { + +GEMMA4_WASM_EXPORT int et_init() { + unload_model(); + release_context(); + try { + g_context = create_webgpu_context(); + } catch (const std::exception& error) { + std::printf("Gemma4 WebGPU initialization failed: %s\n", error.what()); + return 0; + } + if (!compare_and_set_default_webgpu_context(nullptr, &g_context)) { + destroy_webgpu_context(g_context); + std::printf("Gemma4 WebGPU context is already owned\n"); + return 0; + } + g_owns_default_context = true; + return 1; +} + +GEMMA4_WASM_EXPORT int et_load( + const char* pte_path, + const char* data_paths, + const char* method_name) { + unload_model(); + if (!g_owns_default_context || pte_path == nullptr || *pte_path == '\0' || + (method_name != nullptr && *method_name != '\0' && + std::string(method_name) != kMethodName)) { + return 0; + } + auto ptd_paths = parse_ptd_paths(data_paths); + if (ptd_paths.size() != kExpectedPtdCount) { + std::printf( + "Gemma4 requires exactly %zu ordered PTDs, got %zu\n", + kExpectedPtdCount, + ptd_paths.size()); + return 0; + } + + WebGPUModelLoadSpec spec; + spec.pte_path = pte_path; + spec.ptd_paths = std::move(ptd_paths); + spec.required_methods = {kMethodName}; + spec.load_mode = Module::LoadMode::File; + auto next = load_webgpu_model(std::move(spec)); + if (!next.ok() || !validate_loaded_method(*next.get())) { + std::printf("Gemma4 PTE or text_decoder ABI validation failed\n"); + return 0; + } + g_module = std::move(next.get()); + reset_runtime_observations(); + return 1; +} + +GEMMA4_WASM_EXPORT int et_unload() { + unload_model(); + return 1; +} + +GEMMA4_WASM_EXPORT int et_reset() { + if (!g_module) { + return 0; + } + g_module->unload_method(kMethodName); + const Error error = g_module->load_method(kMethodName); + if (error != Error::Ok || !validate_loaded_method(*g_module)) { + unload_model(); + return 0; + } + reset_runtime_observations(); + return 1; +} + +GEMMA4_WASM_EXPORT int et_step(int token, int position) { + return execute_tokens(&token, 1, position); +} + +GEMMA4_WASM_EXPORT void et_prefill_step(int token, int position) { + (void)execute_tokens(&token, 1, position); +} + +GEMMA4_WASM_EXPORT int +et_prefill_batch(const int* tokens, int count, int position) { + if (count == 0 || count == std::numeric_limits::min()) { + return -1; + } + const bool discard_output = count < 0; + const int live_count = discard_output ? -count : count; + const int token = execute_tokens(tokens, live_count, position); + if (token < 0) { + return -1; + } + g_last_prefill_tokens = live_count; + return discard_output ? 0 : token; +} + +GEMMA4_WASM_EXPORT int et_get_last_prefill_token_count() { + return g_last_prefill_tokens; +} + +GEMMA4_WASM_EXPORT int et_get_route_contract_version() { + return 3; +} + +GEMMA4_WASM_EXPORT int et_get_last_route_mask() { +#ifdef WGPU_BACKEND_ENABLE_PROFILING + return static_cast(g_last_route_mask); +#else + return 0; +#endif +} + +GEMMA4_WASM_EXPORT int et_get_last_route_conflict_count() { +#ifdef WGPU_BACKEND_ENABLE_PROFILING + return static_cast(g_last_route_conflict_count); +#else + return 0; +#endif +} + +GEMMA4_WASM_EXPORT void et_profile_enable(int enabled) { +#if defined(_WIN32) + _putenv_s("WEBGPU_TIMESTAMP_QUERY", enabled ? "1" : ""); +#else + if (enabled) { + setenv("WEBGPU_TIMESTAMP_QUERY", "1", 1); + } else { + unsetenv("WEBGPU_TIMESTAMP_QUERY"); + } +#endif +} + +GEMMA4_WASM_EXPORT const char* et_profile() { + static std::string output; +#ifdef WGPU_BACKEND_ENABLE_PROFILING + const auto* context = get_default_webgpu_context(); + if (context == nullptr || !context->querypool || + !context->querypool->results_valid()) { + output = + "{\"perop\":[],\"total_kernel_ms\":0,\"pass_span_ms\":0," + "\"interpass_gap_ms\":0,\"supported\":false}"; + return output.c_str(); + } + + std::map> totals; + double total_ns = 0.0; + uint64_t first_begin = std::numeric_limits::max(); + uint64_t last_end = 0; + for (const auto& duration : context->querypool->results()) { + totals[duration.kernel_name].first += duration.execution_duration_ns / 1e6; + totals[duration.kernel_name].second++; + total_ns += static_cast(duration.execution_duration_ns); + first_begin = std::min(first_begin, duration.start_time_ns); + last_end = std::max(last_end, duration.end_time_ns); + } + const double total_ms = total_ns / 1e6; + const double span_ms = last_end > first_begin + ? static_cast(last_end - first_begin) / 1e6 + : 0.0; + const double gap_ms = std::max(0.0, span_ms - total_ms); + output = "{\"perop\":["; + bool first = true; + for (const auto& entry : totals) { + char item[320]; + const double percent = + total_ms > 0.0 ? 100.0 * entry.second.first / total_ms : 0.0; + std::snprintf( + item, + sizeof(item), + "%s{\"op\":\"%s\",\"ms\":%.4f,\"calls\":%u,\"pct\":%.1f}", + first ? "" : ",", + entry.first.c_str(), + entry.second.first, + entry.second.second, + percent); + output += item; + first = false; + } + char tail[256]; + std::snprintf( + tail, + sizeof(tail), + "],\"total_kernel_ms\":%.4f,\"pass_span_ms\":%.4f," + "\"interpass_gap_ms\":%.4f,\"supported\":true}", + total_ms, + span_ms, + gap_ms); + output += tail; +#else + output = + "{\"perop\":[],\"total_kernel_ms\":0,\"pass_span_ms\":0," + "\"interpass_gap_ms\":0,\"supported\":false}"; +#endif + return output.c_str(); +} + +} // extern "C" + +int main() { + return 0; +} diff --git a/examples/models/gemma4/targets.bzl b/examples/models/gemma4/targets.bzl index fd8179980a7..450c8fb676d 100644 --- a/examples/models/gemma4/targets.bzl +++ b/examples/models/gemma4/targets.bzl @@ -1,4 +1,5 @@ load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") +load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target") GEN_KERNEL_BACKEND_DEPS = [ "//executorch/configurations:optimized_native_cpu_ops", @@ -15,6 +16,37 @@ def _get_torchao_lowbit_deps(): ], }) +def define_webgpu_python_targets(): + fbcode_target(_kind = runtime.python_library, + name = "webgpu_support", + srcs = [ + "webgpu_artifact_manifest.py", + "webgpu_partitioner.py", + ], + _is_external_target = True, + base_module = "executorch.examples.models.gemma4", + resources = { + "config/e2b_config.json": "config/e2b_config.json", + "manifests/gemma4_e2b_webgpu.json": "manifests/gemma4_e2b_webgpu.json", + }, + typing = True, + visibility = ["PUBLIC"], + deps = [ + "//caffe2:torch", + "//executorch/backends/vulkan:op_registry", + "//executorch/backends/vulkan/partitioner:vulkan_partitioner", + "//executorch/backends/vulkan/patterns:vulkan_patterns", + "//executorch/backends/webgpu/scripts:webgpu_artifact_manifest", + "//executorch/exir:lib", + ], + ) + + fbcode_target(_kind = runtime.python_binary, + name = "webgpu_artifact_manifest", + main_function = "executorch.examples.models.gemma4.webgpu_artifact_manifest.main", + deps = [":webgpu_support"], + ) + def define_common_targets(): _KERNEL_BACKEND_DEPS = [ "//executorch/backends/xnnpack:xnnpack_backend", @@ -73,3 +105,15 @@ def define_common_targets(): compiler_flags = ["-Wno-global-constructors"], preprocessor_flags = ["-DET_USE_THREADPOOL"], ) + + runtime.cxx_binary( + name = "gemma4_plain_wasm", + srcs = ["runner/gemma4_plain_wasm.cpp"], + compiler_flags = ["-fexceptions"], + visibility = ["PUBLIC"], + deps = [ + "//executorch/backends/webgpu:webgpu_backend", + "//executorch/backends/webgpu:webgpu_model_loader", + "//executorch/extension/tensor:tensor", + ], + ) diff --git a/examples/models/gemma4/tests/targets.bzl b/examples/models/gemma4/tests/targets.bzl index 92ca1a51ab2..c62ba6d2a64 100644 --- a/examples/models/gemma4/tests/targets.bzl +++ b/examples/models/gemma4/tests/targets.bzl @@ -13,3 +13,18 @@ def define_common_targets(is_fbcode = False): "fbsource//third-party/pypi/transformers:transformers", ], ) + + fbcode_target(_kind = runtime.python_test, + name = "test_webgpu_rewrite_pass", + srcs = ["test_webgpu_rewrite_pass.py"], + preload_deps = [ + "//executorch/extension/llm/custom_ops:custom_ops_aot_lib", + "//executorch/extension/llm/custom_ops:custom_ops_aot_py", + ], + deps = [ + "//caffe2:torch", + "//executorch/backends/vulkan:custom_ops_lib", + "//executorch/examples/models/gemma4:webgpu_support", + "//executorch/exir:lib", + ], + ) diff --git a/examples/models/gemma4/tests/test_webgpu_rewrite_pass.py b/examples/models/gemma4/tests/test_webgpu_rewrite_pass.py new file mode 100644 index 00000000000..0c7f9c45e02 --- /dev/null +++ b/examples/models/gemma4/tests/test_webgpu_rewrite_pass.py @@ -0,0 +1,211 @@ +# 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. + +"""The plain WebGPU edge rewrites must run pre-partition, not in partition(). + +These drive the real `_replace_single_hf_rope` / `_rewrite_gemma4_sdpa` +helpers, not stand-ins, and mutate the exact site counts and ABI guards so a +weakened guard fails. +""" + +import unittest +from typing import List, Optional + +import executorch.backends.vulkan.custom_ops_lib # noqa: F401 +import torch +from executorch.examples.models.gemma4 import webgpu_partitioner as wp +from executorch.exir import to_edge_transform_and_lower +from executorch.exir.backend.compile_spec_schema import CompileSpec +from executorch.exir.backend.partitioner import ( + DelegationSpec, + Partitioner, + PartitionResult, +) +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import ExportedProgramPassBase, ExportedProgramPassResult + + +def _sdpa_available() -> bool: + """llama::custom_sdpa needs its compiled kernel; absent in a bare checkout.""" + try: + exir_ops.edge.llama.custom_sdpa.default + except AttributeError: + return False + return True + + +_REQUIRES_SDPA = unittest.skipUnless( + _sdpa_available(), "llama::custom_sdpa is not registered in this environment" +) + + +class _Stub: + """`_rewrite_gemma4_sdpa` reads only `.graph_module`.""" + + def __init__(self, graph_module: torch.fx.GraphModule) -> None: + self.graph_module = graph_module + + +def _fake(rank: int) -> torch.Tensor: + return torch.zeros([1] * rank) + + +def _sdpa_graph( + count: int, + *, + args: int = 8, + mask_rank: int = 2, + dropout: float = 0.0, + is_causal: bool = False, + scale: float = 1.0, +) -> torch.fx.GraphModule: + graph = torch.fx.Graph() + ph = {r: graph.placeholder(f"p{r}") for r in (2, 4)} + for r, node in ph.items(): + node.meta["val"] = _fake(r) + for i in range(count): + full = ( + ph[4], + ph[4], + ph[4], + 0, + ph[mask_rank], + dropout, + is_causal, + scale, + ) + node = graph.call_function(exir_ops.edge.llama.custom_sdpa.default, full[:args]) + node.meta["val"] = _fake(4) + node.name = f"sdpa_{i}" + graph.output(ph[4]) + return torch.fx.GraphModule(torch.nn.Module(), graph) + + +class _Inert(Partitioner): + def __init__(self) -> None: + super().__init__() + self.delegation_spec = DelegationSpec("VulkanBackend", [CompileSpec("k", b"v")]) + self.seen: List[object] = [] + self.input_ids: Optional[List[int]] = None + + def partition(self, exported_program) -> PartitionResult: + nodes = list(exported_program.graph_module.graph.nodes) + self.seen = [n.target for n in nodes if n.op == "call_function"] + self.input_ids = [id(n) for n in nodes] + return PartitionResult( + tagged_exported_program=exported_program, partition_tags={} + ) + + +class _Probe(ExportedProgramPassBase): + def __init__(self, log: List[str], name: str) -> None: + super().__init__() + self.log, self.name = log, name + + def call(self, ep) -> ExportedProgramPassResult: + self.log.append(self.name) + return ExportedProgramPassResult(ep, True) + + +class RewritePassTest(unittest.TestCase): + @_REQUIRES_SDPA + def test_exact_sdpa_count_is_enforced(self) -> None: + # The accepted contract is exactly 35 sites. + wp._rewrite_gemma4_sdpa(_Stub(_sdpa_graph(35))) + for wrong in (0, 34, 36): + with self.subTest(count=wrong): + with self.assertRaisesRegex(ValueError, "35 SDPA sites"): + wp._rewrite_gemma4_sdpa(_Stub(_sdpa_graph(wrong))) + + @_REQUIRES_SDPA + def test_rewrite_retargets_every_site(self) -> None: + gm = _sdpa_graph(35) + wp._rewrite_gemma4_sdpa(_Stub(gm)) + targets = [n.target for n in gm.graph.nodes if n.op == "call_function"] + self.assertEqual(targets.count(exir_ops.edge.et_vk.gemma4_sdpa.default), 35) + self.assertNotIn(exir_ops.edge.llama.custom_sdpa.default, targets) + + @_REQUIRES_SDPA + def test_abi_and_rank_guards_fail_closed(self) -> None: + cases = { + "positional ABI": dict(args=7), + "not WebGPU-compatible": dict(mask_rank=4), + } + for expected, kwargs in cases.items(): + with self.subTest(**kwargs): + with self.assertRaisesRegex(ValueError, expected): + wp._rewrite_gemma4_sdpa(_Stub(_sdpa_graph(35, **kwargs))) + for kwargs in ( + dict(dropout=0.1), + dict(is_causal=True), + dict(scale=0.5), + ): + with self.subTest(**kwargs): + with self.assertRaisesRegex(ValueError, "not WebGPU-compatible"): + wp._rewrite_gemma4_sdpa(_Stub(_sdpa_graph(35, **kwargs))) + + def test_rope_count_is_enforced(self) -> None: + # No RoPE sites in an empty graph: the 20-site contract must fire. + empty = torch.fx.GraphModule(torch.nn.Module(), torch.fx.Graph()) + with self.assertRaisesRegex(ValueError, "20 single-HF-RoPE sites"): + wp._replace_single_hf_rope(_Stub(empty)) + + def test_pass_applies_both_rewrites_in_order(self) -> None: + # RoPE runs first, so its contract is what an empty graph trips. + with self.assertRaisesRegex(ValueError, "single-HF-RoPE"): + wp._Gemma4WebGPURewritePass().call( + _Stub(torch.fx.GraphModule(torch.nn.Module(), torch.fx.Graph())) + ) + + def test_transform_passes_entry_point_is_fixed(self) -> None: + passes = wp.build_webgpu_transform_passes() + self.assertEqual(len(passes), 1) + self.assertIsInstance(passes[0], wp._Gemma4WebGPURewritePass) + self.assertFalse(hasattr(passes[0], "rewrites")) + + +class _Add(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return torch.ops.aten.add.Tensor(x, x) + + +class _Mutating(Partitioner): + def __init__(self) -> None: + super().__init__() + self.delegation_spec = DelegationSpec("VulkanBackend", [CompileSpec("k", b"v")]) + + def partition(self, exported_program) -> PartitionResult: + for node in exported_program.graph_module.graph.nodes: + if node.target == exir_ops.edge.aten.add.Tensor: + node.target = exir_ops.edge.aten.mul.Tensor + exported_program.graph_module.recompile() + return PartitionResult( + tagged_exported_program=exported_program, partition_tags={} + ) + + +class PlacementTest(unittest.TestCase): + def _program(self): + return torch.export.export(_Add(), (torch.randn(4),), strict=True) + + def test_partition_time_mutation_is_rejected(self) -> None: + with self.assertRaises(AssertionError) as caught: + to_edge_transform_and_lower(self._program(), partitioner=[_Mutating()]) + self.assertIn("should not modify the graph module", str(caught.exception)) + + def test_partition_input_is_left_identical(self) -> None: + inert = _Inert() + to_edge_transform_and_lower(self._program(), partitioner=[inert]) + self.assertIn(exir_ops.edge.aten.add.Tensor, inert.seen) + + def test_transform_ordering_is_deterministic(self) -> None: + log: List[str] = [] + to_edge_transform_and_lower( + self._program(), + partitioner=[_Inert()], + transform_passes=[_Probe(log, "a"), _Probe(log, "b")], + ) + self.assertEqual(log, ["a", "b"]) diff --git a/examples/models/gemma4/text_decoder/gemma4_attention.py b/examples/models/gemma4/text_decoder/gemma4_attention.py index 6d4550013dc..88106f9afd5 100644 --- a/examples/models/gemma4/text_decoder/gemma4_attention.py +++ b/examples/models/gemma4/text_decoder/gemma4_attention.py @@ -34,6 +34,30 @@ def rotate_half(x: torch.Tensor) -> torch.Tensor: return torch.cat((-x2, x1), dim=-1) +def precompute_freqs_cis( + head_dim: int, + rotary_dim: int, + end: int, + theta: float, +) -> Tuple[torch.Tensor, torch.Tensor]: + if head_dim <= 0 or rotary_dim <= 0 or rotary_dim > head_dim: + raise ValueError("RoPE dimensions must satisfy 0 < rotary_dim <= head_dim") + if head_dim % 2 != 0 or rotary_dim % 2 != 0: + raise ValueError("RoPE dimensions must be even") + + inv_freq_rotated = 1.0 / ( + theta ** (torch.arange(0, rotary_dim, 2).float() / head_dim) + ) + nope_angles = head_dim // 2 - rotary_dim // 2 + inv_freq = torch.cat( + (inv_freq_rotated, torch.zeros(nope_angles, dtype=torch.float32)) + ) + positions = torch.arange(end, dtype=torch.float32) + freqs = torch.outer(positions, inv_freq) + emb = torch.cat((freqs, freqs), dim=-1) + return torch.cos(emb), torch.sin(emb) + + def apply_rotary_emb( xq: torch.Tensor, xk: torch.Tensor, @@ -43,16 +67,16 @@ def apply_rotary_emb( """Apply rotary position embeddings to Q and K (HuggingFace style). Args: - xq: Query tensor of shape [batch, num_heads, seq_len, rotary_dim] - xk: Key tensor of shape [batch, num_kv_heads, seq_len, rotary_dim] + xq: Query tensor of shape [batch, seq_len, num_heads, head_dim] + xk: Key tensor of shape [batch, seq_len, num_kv_heads, head_dim] freqs_cos: Cosine frequencies [seq_len, rotary_dim] freqs_sin: Sine frequencies [seq_len, rotary_dim] Returns: Tuple of (rotated_q, rotated_k) """ - freqs_cos = freqs_cos.unsqueeze(0).unsqueeze(0) - freqs_sin = freqs_sin.unsqueeze(0).unsqueeze(0) + freqs_cos = freqs_cos.unsqueeze(1) + freqs_sin = freqs_sin.unsqueeze(1) xq_out = (xq.float() * freqs_cos) + (rotate_half(xq.float()) * freqs_sin) xk_out = (xk.float() * freqs_cos) + (rotate_half(xk.float()) * freqs_sin) @@ -66,8 +90,8 @@ def apply_rotary_emb_single( freqs_sin: torch.Tensor, ) -> torch.Tensor: """Apply rotary position embeddings to a single tensor (Q only).""" - freqs_cos = freqs_cos.unsqueeze(0).unsqueeze(0) - freqs_sin = freqs_sin.unsqueeze(0).unsqueeze(0) + freqs_cos = freqs_cos.unsqueeze(1) + freqs_sin = freqs_sin.unsqueeze(1) x_out = (x.float() * freqs_cos) + (rotate_half(x.float()) * freqs_sin) @@ -95,8 +119,10 @@ def __init__( use_index_copy: bool = False, ): super().__init__() + from executorch.extension.llm.custom_ops import custom_ops # noqa: F401 + self.use_index_copy = use_index_copy - cache_shape = (max_batch_size, num_kv_heads, max_seq_len, head_dim) + cache_shape = (max_batch_size, max_seq_len, num_kv_heads, head_dim) self.register_buffer( "k_cache", torch.zeros(cache_shape, dtype=dtype), persistent=False ) @@ -116,15 +142,26 @@ def update( Tuple of (full_k, full_v) - returns entire cache """ if self.use_index_copy: - self.k_cache.index_copy_(2, input_pos, k_val) - self.v_cache.index_copy_(2, input_pos, v_val) + self.k_cache.index_copy_(1, input_pos, k_val) + self.v_cache.index_copy_(1, input_pos, v_val) else: - seq_len = k_val.size(2) + seq_len = k_val.size(1) start_pos = input_pos[0].item() torch._check_is_size(start_pos) torch._check(start_pos >= 0) - self.k_cache.narrow(2, start_pos, seq_len).copy_(k_val) - self.v_cache.narrow(2, start_pos, seq_len).copy_(v_val) + if not torch.compiler.is_compiling(): + if start_pos + seq_len > self.k_cache.size(1): + raise ValueError("Gemma4 KV update exceeds cache capacity") + expected_pos = torch.arange( + start_pos, + start_pos + seq_len, + dtype=input_pos.dtype, + device=input_pos.device, + ) + if input_pos.dim() != 1 or not torch.equal(input_pos, expected_pos): + raise ValueError("Gemma4 KV positions must be contiguous") + torch.ops.llama.update_cache.default(k_val, self.k_cache, start_pos) + torch.ops.llama.update_cache.default(v_val, self.v_cache, start_pos) return self.k_cache, self.v_cache @@ -201,19 +238,8 @@ def __init__( else: self.rotary_dim = self.head_dim - # RoPE: store only inv_freq; cos/sin computed on the fly per forward. - # Partial RoPE pads with zeros for non-rotated dims so rotate_half pairs correctly. - rope_angles = self.rotary_dim // 2 - inv_freq_rotated = 1.0 / ( - self.rope_theta - ** (torch.arange(0, self.rotary_dim, 2).float() / self.head_dim) - ) - nope_angles = self.head_dim // 2 - rope_angles - if nope_angles > 0: - inv_freq = torch.cat([inv_freq_rotated, torch.zeros(nope_angles)]) - else: - inv_freq = inv_freq_rotated - self.register_buffer("inv_freq", inv_freq, persistent=False) + self.register_buffer("freqs_cos", None, persistent=False) + self.register_buffer("freqs_sin", None, persistent=False) # KV cache — skip allocation for shared layers (they use donor's KV) self.use_index_copy = config.use_index_copy_for_kv_cache @@ -269,12 +295,32 @@ def _get_rope_freqs( self, input_pos: Optional[torch.Tensor], seq_len: int, + query_start_pos: Optional[int] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: - """Compute RoPE cos/sin from inv_freq for the current positions.""" - pos = input_pos if input_pos is not None else torch.arange(seq_len) - freqs = torch.outer(pos.float(), self.inv_freq) - emb = torch.cat((freqs, freqs), dim=-1) - return torch.cos(emb), torch.sin(emb) + """Slice shared RoPE tables for the current contiguous positions.""" + if self.freqs_cos is None or self.freqs_sin is None: + raise RuntimeError("Gemma4 RoPE tables were not initialized") + if ( + self.use_index_copy + and query_start_pos is None + and input_pos is not None + ): + return ( + torch.index_select(self.freqs_cos, 0, input_pos), + torch.index_select(self.freqs_sin, 0, input_pos), + ) + start_pos = ( + query_start_pos + if query_start_pos is not None + else (0 if input_pos is None else input_pos[0].item()) + ) + torch._check_is_size(start_pos) + torch._check(start_pos >= 0) + torch._check(start_pos + seq_len <= self.freqs_cos.size(0)) + return ( + self.freqs_cos.narrow(0, start_pos, seq_len), + self.freqs_sin.narrow(0, start_pos, seq_len), + ) def _slice_cache_to_actual( self, @@ -300,9 +346,9 @@ def _slice_cache_to_actual( actual_kv_len = start_pos + seq_len torch._check_is_size(actual_kv_len) torch._check(actual_kv_len >= seq_len) - torch._check(actual_kv_len <= k.size(2)) - k = k.narrow(2, 0, actual_kv_len) - v = v.narrow(2, 0, actual_kv_len) + torch._check(actual_kv_len <= k.size(1)) + k = k.narrow(1, 0, actual_kv_len) + v = v.narrow(1, 0, actual_kv_len) if attn_mask is not None: attn_mask = attn_mask.narrow(1, 0, actual_kv_len) return k, v, attn_mask @@ -313,8 +359,15 @@ def _slice_mask( input_pos: torch.Tensor, seq_len: int, kv_len: int, + query_start_pos: Optional[int] = None, ) -> torch.Tensor: """Slice a [max_seq_len, max_seq_len] mask to current query positions x cache.""" + if query_start_pos is not None: + torch._check_is_size(query_start_pos) + torch._check(query_start_pos >= 0) + return base_mask.narrow(0, query_start_pos, seq_len).narrow( + 1, 0, kv_len + ) if self.use_index_copy: return torch.index_select(base_mask, 0, input_pos).narrow(1, 0, kv_len) start_pos = input_pos[0].item() @@ -327,6 +380,7 @@ def _build_attn_mask( input_pos: Optional[torch.Tensor], seq_len: int, kv_len: int, + query_start_pos: Optional[int] = None, ) -> torch.Tensor: """Combined causal + sliding-window mask for the current step.""" using_cached_kv = ( @@ -335,14 +389,20 @@ def _build_attn_mask( and kv_len > seq_len ) if using_cached_kv: - mask = self._slice_mask(self.causal_mask, input_pos, seq_len, kv_len) + mask = self._slice_mask( + self.causal_mask, input_pos, seq_len, kv_len, query_start_pos + ) else: mask = self.causal_mask[:seq_len, :seq_len] if self.sliding_window is not None and self.sliding_window_mask is not None: if using_cached_kv: sw_mask = self._slice_mask( - self.sliding_window_mask, input_pos, seq_len, kv_len + self.sliding_window_mask, + input_pos, + seq_len, + kv_len, + query_start_pos, ) else: sw_mask = self.sliding_window_mask[:seq_len, :seq_len] @@ -355,6 +415,7 @@ def forward( input_pos: Optional[torch.Tensor] = None, mask: Optional[torch.Tensor] = None, shared_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + query_start_pos: Optional[int] = None, ) -> Tuple[torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]: """Forward pass for attention. @@ -363,6 +424,7 @@ def forward( input_pos: Current position(s) for KV cache update mask: Optional attention mask shared_kv: Optional tuple of (k, v) from donor layer for YOCO + query_start_pos: Optional absolute position for a selected query row Returns: Tuple of: @@ -373,32 +435,32 @@ def forward( # Compute Q projection q = self.q_proj(hidden_states) - q = q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2) + q = q.view(batch_size, seq_len, self.num_heads, self.head_dim) q = self.q_norm(q) # For KV shared layers, use shared K/V from donor layer if self.is_kv_shared_layer and shared_kv is not None: k, v = shared_kv - freqs_cos, freqs_sin = self._get_rope_freqs(input_pos, seq_len) + freqs_cos, freqs_sin = self._get_rope_freqs( + input_pos, seq_len, query_start_pos + ) q = self._apply_rope_single(q, freqs_cos, freqs_sin) else: # Compute K, V projections k = self.k_proj(hidden_states) v = self.v_proj(hidden_states) - k = k.view(batch_size, seq_len, self.num_kv_heads, self.head_dim).transpose( - 1, 2 - ) - v = v.view(batch_size, seq_len, self.num_kv_heads, self.head_dim).transpose( - 1, 2 - ) + k = k.view(batch_size, seq_len, self.num_kv_heads, self.head_dim) + v = v.view(batch_size, seq_len, self.num_kv_heads, self.head_dim) # Apply KV norms k = self.k_norm(k) v = self.v_norm(v) # Get RoPE frequencies - freqs_cos, freqs_sin = self._get_rope_freqs(input_pos, seq_len) + freqs_cos, freqs_sin = self._get_rope_freqs( + input_pos, seq_len, query_start_pos + ) # Apply RoPE (partial for full attention, full for sliding) q, k = self._apply_rope(q, k, freqs_cos, freqs_sin) @@ -430,53 +492,66 @@ def forward( # and tiles attention so the [seq x seq] matrix never materializes. # Build mask with the static kv_len first (so _build_attn_mask sees # a constant), then slice K/V/mask to the actually-valid range. - full_kv_len = k.size(2) - start_pos = 0 if self.use_index_copy else input_pos[0].item() - attn_mask = self._build_attn_mask(input_pos, seq_len, full_kv_len) + full_kv_len = k.size(1) + start_pos = ( + query_start_pos + if query_start_pos is not None + else (0 if self.use_index_copy else input_pos[0].item()) + ) + attn_mask = self._build_attn_mask( + input_pos, seq_len, full_kv_len, query_start_pos + ) if not self.use_index_copy: k, v, attn_mask = self._slice_cache_to_actual( k, v, start_pos, seq_len, attn_mask ) - # custom_sdpa expects [bs, seq_len, n_heads, head_dim] - q_sdpa = q.transpose(1, 2) - k_sdpa = k.transpose(1, 2) - v_sdpa = v.transpose(1, 2) - # custom_sdpa positional args: (q, k, v, start_pos, attn_mask, dropout, is_causal, scale). # The op schema has a typo (`drpout_p`); avoid kwargs. attn_output = torch.ops.llama.custom_sdpa( - q_sdpa, - k_sdpa, - v_sdpa, + q, + k, + v, start_pos, attn_mask, 0.0, False, self.scaling, ) - attn_output = attn_output.view(batch_size, seq_len, -1) + attn_output = attn_output.view( + batch_size, seq_len, self.num_heads * self.head_dim + ) else: # Same cache-slice optimization as the custom_sdpa branch above. # Build mask first with full kv_len (so _build_attn_mask sees a # constant), then slice K/V/mask to the actually-valid range. if mask is None: - mask = self._build_attn_mask(input_pos, seq_len, k.size(2)) + mask = self._build_attn_mask( + input_pos, seq_len, k.size(1), query_start_pos + ) if not self.use_index_copy and input_pos is not None: + start_pos = ( + query_start_pos + if query_start_pos is not None + else input_pos[0].item() + ) k, v, mask = self._slice_cache_to_actual( - k, v, input_pos[0].item(), seq_len, mask + k, v, start_pos, seq_len, mask ) - k = self._repeat_kv(k) - v = self._repeat_kv(v) - attn_weights = torch.matmul(q, k.transpose(-2, -1)) * self.scaling + q_bhsd = q.transpose(1, 2) + k_bhsd = self._repeat_kv(k.transpose(1, 2)) + v_bhsd = self._repeat_kv(v.transpose(1, 2)) + attn_weights = torch.matmul(q_bhsd, k_bhsd.transpose(-2, -1)) * self.scaling attn_weights = attn_weights + mask.unsqueeze(0).unsqueeze(0) attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32).type_as( - q + q_bhsd ) - attn_output = torch.matmul(attn_weights, v) + attn_output = torch.matmul(attn_weights, v_bhsd) attn_output = attn_output.transpose(1, 2).contiguous() - attn_output = attn_output.view(batch_size, seq_len, -1) + attn_output = attn_output.view( + batch_size, seq_len, self.num_heads * self.head_dim + ) attn_output = self.o_proj(attn_output) @@ -513,8 +588,8 @@ def __init__( self.use_index_copy = use_index_copy self.return_float_values = return_float_values - cache_shape = (max_batch_size, num_kv_heads, max_seq_len, head_dim) - scale_shape = (max_batch_size, num_kv_heads, max_seq_len, 1) + cache_shape = (max_batch_size, max_seq_len, num_kv_heads, head_dim) + scale_shape = (max_batch_size, max_seq_len, num_kv_heads, 1) self.register_buffer("k_cache", torch.zeros(cache_shape, dtype=torch.int8)) self.register_buffer("v_cache", torch.zeros(cache_shape, dtype=torch.int8)) @@ -549,19 +624,19 @@ def update( quantized_v, v_scales = self._quantize(v_val) if self.use_index_copy: - self.k_cache.index_copy_(2, input_pos, quantized_k) - self.v_cache.index_copy_(2, input_pos, quantized_v) - self.k_cache_scales.index_copy_(2, input_pos, k_scales) - self.v_cache_scales.index_copy_(2, input_pos, v_scales) + self.k_cache.index_copy_(1, input_pos, quantized_k) + self.v_cache.index_copy_(1, input_pos, quantized_v) + self.k_cache_scales.index_copy_(1, input_pos, k_scales) + self.v_cache_scales.index_copy_(1, input_pos, v_scales) else: - seq_len = k_val.size(2) + seq_len = k_val.size(1) start_pos = input_pos[0].item() torch._check_is_size(start_pos) torch._check(start_pos >= 0) - self.k_cache.narrow(2, start_pos, seq_len).copy_(quantized_k) - self.v_cache.narrow(2, start_pos, seq_len).copy_(quantized_v) - self.k_cache_scales.narrow(2, start_pos, seq_len).copy_(k_scales) - self.v_cache_scales.narrow(2, start_pos, seq_len).copy_(v_scales) + self.k_cache.narrow(1, start_pos, seq_len).copy_(quantized_k) + self.v_cache.narrow(1, start_pos, seq_len).copy_(quantized_v) + self.k_cache_scales.narrow(1, start_pos, seq_len).copy_(k_scales) + self.v_cache_scales.narrow(1, start_pos, seq_len).copy_(v_scales) if not self.return_float_values: return self.k_cache, self.v_cache @@ -571,11 +646,11 @@ def update( v_out = (self.v_cache.to(torch.float32) * self.v_cache_scales).to(self.dtype) if self.use_index_copy: - k_out.index_copy_(2, input_pos, k_val) - v_out.index_copy_(2, input_pos, v_val) + k_out.index_copy_(1, input_pos, k_val) + v_out.index_copy_(1, input_pos, v_val) else: - k_out.narrow(2, start_pos, seq_len).copy_(k_val) - v_out.narrow(2, start_pos, seq_len).copy_(v_val) + k_out.narrow(1, start_pos, seq_len).copy_(k_val) + v_out.narrow(1, start_pos, seq_len).copy_(v_val) return k_out, v_out @@ -584,7 +659,7 @@ def from_float( cls, kv_cache: Gemma4KVCache, return_float_values: bool = True ) -> "Gemma4QuantizedKVCache": """Create quantized KV cache from float KV cache.""" - max_batch_size, num_kv_heads, max_seq_len, head_dim = kv_cache.k_cache.shape + max_batch_size, max_seq_len, num_kv_heads, head_dim = kv_cache.k_cache.shape dtype = kv_cache.k_cache.dtype return cls( max_batch_size, diff --git a/examples/models/gemma4/text_decoder/gemma4_cross_decoder.py b/examples/models/gemma4/text_decoder/gemma4_cross_decoder.py index af1f77a4742..5c91a499526 100644 --- a/examples/models/gemma4/text_decoder/gemma4_cross_decoder.py +++ b/examples/models/gemma4/text_decoder/gemma4_cross_decoder.py @@ -52,6 +52,7 @@ def forward( per_layer_inputs: torch.Tensor, shared_kv: Dict[int, Tuple[torch.Tensor, torch.Tensor]], input_pos: Optional[torch.Tensor] = None, + query_start_pos: Optional[int] = None, ) -> torch.Tensor: """Forward pass through cross-decoder. @@ -60,6 +61,7 @@ def forward( per_layer_inputs: Remaining per-layer inputs from self-decoder shared_kv: Dict mapping donor layer indices to (k, v) tuples input_pos: Current position(s) for KV cache + query_start_pos: Optional absolute position for a selected query row Returns: hidden_states: [batch, seq_len, hidden_size] @@ -78,6 +80,7 @@ def forward( per_layer_input=per_layer_input, input_pos=input_pos, shared_kv=layer_shared_kv, + query_start_pos=query_start_pos, ) return hidden_states diff --git a/examples/models/gemma4/text_decoder/gemma4_decoder_layer.py b/examples/models/gemma4/text_decoder/gemma4_decoder_layer.py index e10c1c7e415..822fd3caea3 100644 --- a/examples/models/gemma4/text_decoder/gemma4_decoder_layer.py +++ b/examples/models/gemma4/text_decoder/gemma4_decoder_layer.py @@ -126,6 +126,7 @@ def forward( per_layer_input: torch.Tensor, input_pos: Optional[torch.Tensor] = None, shared_kv: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, + query_start_pos: Optional[int] = None, ) -> Tuple[torch.Tensor, torch.Tensor, Optional[Tuple[torch.Tensor, torch.Tensor]]]: """Forward pass for decoder layer. @@ -134,6 +135,7 @@ def forward( per_layer_input: Per-layer input of shape [batch, seq_len, hidden_size_per_layer_input] input_pos: Current position(s) for KV cache shared_kv: Optional tuple of (k, v) from donor layer for YOCO + query_start_pos: Optional absolute position for a selected query row Returns: Tuple of (hidden_states, per_layer_input, kv_to_share) @@ -145,6 +147,7 @@ def forward( hidden_states=hidden_states, input_pos=input_pos, shared_kv=shared_kv, + query_start_pos=query_start_pos, ) hidden_states = self.post_attention_layernorm(hidden_states) hidden_states = residual + hidden_states diff --git a/examples/models/gemma4/text_decoder/gemma4_model.py b/examples/models/gemma4/text_decoder/gemma4_model.py index 3c2b985507d..5dcdac63d3f 100644 --- a/examples/models/gemma4/text_decoder/gemma4_model.py +++ b/examples/models/gemma4/text_decoder/gemma4_model.py @@ -153,7 +153,9 @@ def get_example_inputs_with_audio( return (input_ids, inputs_embeds) def get_dynamic_shapes( - self, with_audio_embeds: bool = False + self, + with_audio_embeds: bool = False, + max_input_len: Optional[int] = None, ) -> Optional[Dict[str, Any]]: """Get dynamic shape specifications for export.""" if not self.config.enable_dynamic_shape: @@ -161,7 +163,16 @@ def get_dynamic_shapes( from torch.export import Dim - seq_len = Dim("seq_len", min=1, max=self.config.max_seq_len - 1) + input_limit = ( + self.config.max_seq_len - 1 + if max_input_len is None + else max_input_len + ) + if input_limit < 2 or input_limit >= self.config.max_seq_len: + raise ValueError( + "max_input_len must be at least 2 and less than max_seq_len" + ) + seq_len = Dim("seq_len", min=1, max=input_limit) if self.config.use_kv_cache: if with_audio_embeds: diff --git a/examples/models/gemma4/text_decoder/gemma4_self_decoder.py b/examples/models/gemma4/text_decoder/gemma4_self_decoder.py index 2cc00b838f7..53926eb3b81 100644 --- a/examples/models/gemma4/text_decoder/gemma4_self_decoder.py +++ b/examples/models/gemma4/text_decoder/gemma4_self_decoder.py @@ -119,12 +119,15 @@ def _compute_per_layer_inputs( per_layer_proj = self.per_layer_projection_norm(per_layer_proj) # Get per-layer token embeddings - per_layer_mask = torch.logical_and( - input_ids >= 0, input_ids < self.config.vocab_size_per_layer_input - ) - per_layer_tokens = torch.where( - per_layer_mask, input_ids, torch.zeros_like(input_ids) - ) + if self.config.vocab_size == self.config.vocab_size_per_layer_input: + per_layer_tokens = input_ids + else: + per_layer_mask = torch.logical_and( + input_ids >= 0, input_ids < self.config.vocab_size_per_layer_input + ) + per_layer_tokens = torch.where( + per_layer_mask, input_ids, torch.zeros_like(input_ids) + ) per_layer_embed = ( self.embed_tokens_per_layer(per_layer_tokens) * self.embed_scale_per_layer ) diff --git a/examples/models/gemma4/text_decoder/gemma4_transformer.py b/examples/models/gemma4/text_decoder/gemma4_transformer.py index ad078f0dd4d..b29f8ea0636 100644 --- a/examples/models/gemma4/text_decoder/gemma4_transformer.py +++ b/examples/models/gemma4/text_decoder/gemma4_transformer.py @@ -20,6 +20,7 @@ import torch from torch import nn +from .gemma4_attention import Gemma4Attention, precompute_freqs_cis from .gemma4_config import Gemma4Config from .gemma4_cross_decoder import Gemma4CrossDecoder from .gemma4_self_decoder import Gemma4SelfDecoder @@ -70,6 +71,7 @@ def __init__(self, config: Gemma4Config): # Create shared masks (one copy each, referenced by all attention layers) self._share_masks(config) + self._share_rope_tables(config) def _share_masks(self, config: Gemma4Config) -> None: """Create causal and sliding window masks once, share across all attention layers. @@ -100,6 +102,26 @@ def _share_masks(self, config: Gemma4Config) -> None: ): module.sliding_window_mask = sw_mask + def _share_rope_tables(self, config: Gemma4Config) -> None: + local_tables = precompute_freqs_cis( + config.head_dim, + config.head_dim, + config.max_seq_len, + config.rope_local_base_freq, + ) + global_tables = precompute_freqs_cis( + config.global_head_dim, + int(config.global_head_dim * config.partial_rotary_factor), + config.max_seq_len, + config.rope_theta, + ) + + for module in self.modules(): + if not isinstance(module, Gemma4Attention): + continue + tables = local_tables if module.is_sliding else global_tables + module.freqs_cos, module.freqs_sin = tables + def forward( self, input_ids: torch.Tensor, @@ -123,19 +145,26 @@ def forward( inputs_embeds=inputs_embeds, ) + query_start_pos = None + if input_pos is not None: + query_start_pos = input_pos[-1].item() + hidden_states = hidden_states[:, -1:, :] + per_layer_inputs = per_layer_inputs[:, :, -1:, :] + # Cross-decoder (with shared K/V from self-decoder for YOCO) hidden_states = self.cross_decoder( hidden_states=hidden_states, per_layer_inputs=per_layer_inputs, shared_kv=shared_kv, input_pos=input_pos, + query_start_pos=query_start_pos, ) # Final normalization hidden_states = self.norm(hidden_states) - # Only compute lm_head for the last token during prefill - hidden_states = hidden_states[:, -1:, :] + if input_pos is None: + hidden_states = hidden_states[:, -1:, :] # Output projection logits = self.lm_head(hidden_states) diff --git a/examples/models/gemma4/webgpu_artifact_manifest.py b/examples/models/gemma4/webgpu_artifact_manifest.py new file mode 100644 index 00000000000..8644d80de21 --- /dev/null +++ b/examples/models/gemma4/webgpu_artifact_manifest.py @@ -0,0 +1,911 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-strict + +"""Gemma 4 E2B acquisition and WebGPU artifact manifest contract.""" + +import argparse +import copy +import hashlib +import importlib.util +import json +import subprocess + +from pathlib import Path +from typing import Any, Mapping, Sequence + +from executorch.backends.webgpu.scripts.webgpu_artifact_manifest import ( + create_manifest, + validate_manifest, +) + + +SOURCE_CONFIG_SHA256 = ( + "526e9fd34a8a489c35952535335a4b8556e9169d851187a895f80286e7466206" +) +SOURCE_CONFIG_BYTES = 2214 +WEBGPU_EXTERNAL_CONSTANTS_MAX_DATA_BYTES = 1_500_000_000 +EXPECTED_LAYER_TYPES: list[str] = [ + attention_type + for _ in range(7) + for attention_type in [ + "sliding_attention", + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + ] +] +ARCHITECTURE_FINGERPRINT: dict[str, object] = { + "model_type": "gemma4", + "num_hidden_layers": 35, + "hidden_size": 1536, + "intermediate_size": 6144, + "num_attention_heads": 8, + "num_key_value_heads": 1, + "head_dim": 256, + "global_head_dim": 512, + "num_kv_shared_layers": 20, + "vocab_size": 262144, + "layer_types": EXPECTED_LAYER_TYPES, +} +CHECKPOINT_ACQUISITION: dict[str, object] = { + "repo_id": "google/gemma-4-E2B-it-qat-q4_0-unquantized", + "revision": "6befbaca7398925921802abd1f277b495b78b738", + "files": { + "model.safetensors": { + "bytes": 10208852878, + "sha256": "33fe0cece08fb527ffefbd1a3a9ce73bd71073727993a283506293e5c6bf0137", + }, + "config.json": { + "bytes": 4946, + "sha256": "bbeff1e2fd3fe282536e7ace02309d43e0dbd9b6ac4b6a149b97e3ab6942a878", + }, + "tokenizer.json": { + "bytes": 32169626, + "sha256": "cc8d3a0ce36466ccc1278bf987df5f71db1719b9ca6b4118264f45cb627bfe0f", + }, + "tokenizer_config.json": { + "bytes": 3729, + "sha256": "3ab5c7b94dc97d65ca7064496fa69b88ff875378e1cb7ee3e43070c3a8170999", + }, + "generation_config.json": { + "bytes": 203, + "sha256": "b69207f9be617e982d13cc273cce6fd88c98dda99a4bdc5e2d52ffe0a0d9f0a9", + }, + "processor_config.json": { + "bytes": 1689, + "sha256": "32bdf45d2ad4cc29a0822ddd157a182de76644f0419a6228d151495256e9813c", + }, + "chat_template.jinja": { + "bytes": 18569, + "sha256": "0a2c8073c878ab1da004bee933a998606537bbb62016310352c7285c3f01c5b5", + }, + "README.md": { + "bytes": 29351, + "sha256": "aaab87052837925e0fb400bb20700553b11088fa5b3ae21fa0c1ec5da53637a4", + }, + ".gitattributes": { + "bytes": 1570, + "sha256": "34448b82c17d60fec9b65b1f093c115ddbaadc04beb1b0140b6bfed2e012a930", + }, + }, +} +EXPORT_CONTRACT: dict[str, object] = { + "backend": "webgpu", + "max_input_len": 512, + "max_seq_len": 8960, + "methods": ["text_decoder"], + "output": { + "dtype": "int64", + "semantic": "greedy_token", + "shape": [1, 1], + }, + "quantization": "8da4w+emb4", +} +_SOURCE_CLOSURE_RECEIPT_SCHEMA_VERSION = 3 +_SOURCE_MANIFEST_SCHEMA_VERSION = 1 +_WGSL_MANIFEST_SCHEMA_VERSION = 1 +_GEMMA_PRODUCTION_DIFF_SUMMARIES: tuple[str, ...] = ( + "[ExecuTorch][WebGPU] Add shared model runtime prerequisites", + "[ExecuTorch][Vulkan] Support scoped Gemma symbolic partitioning", + "[ExecuTorch][WebGPU] Add Gemma 4 plain runtime and guarded routes", + "[ExecuTorch][WebGPU] Add Gemma 4 plain export and artifact contract", + "[ExecuTorch][WebGPU] Add plain Gemma 4 source-closure tests", +) + + +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]: + return { + "schema_version": 1, + "artifacts": [ + { + "bytes": byte_count, + "path": path, + "role": "source", + "sha256": sha256, + } + ], + "ptd_order": [], + } + + +def _load_json(path: Path) -> Mapping[str, object]: + value = json.loads(path.read_text(encoding="utf-8")) + if not isinstance(value, dict): + raise ValueError(f"JSON document must be an object: {path}") + return value + + +def _validate_architecture(config: Mapping[str, object], label: str) -> None: + text_config = config.get("text_config") + if not isinstance(text_config, dict): + raise ValueError(f"{label} is missing text_config") + observed: dict[str, object] = {"model_type": config.get("model_type")} + for key in ARCHITECTURE_FINGERPRINT: + if key != "model_type": + observed[key] = text_config.get(key) + if observed != ARCHITECTURE_FINGERPRINT: + raise ValueError(f"{label} architecture fingerprint mismatch") + + +def _is_hex_digest(value: object, length: int) -> bool: + return ( + isinstance(value, str) + and len(value) == length + and all(character in "0123456789abcdef" for character in value) + ) + + +def _sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as source: + for block in iter(lambda: source.read(1024 * 1024), b""): + digest.update(block) + return digest.hexdigest() + + +def _require_exact_keys( + value: Mapping[str, object], expected: set[str], label: str +) -> None: + if set(value) != expected: + raise ValueError(f"{label} schema mismatch: {sorted(value)}") + + +def _canonical_set_digest(value: object) -> str: + encoded = json.dumps(value, separators=(",", ":"), sort_keys=True).encode("utf-8") + return hashlib.sha256(encoded).hexdigest() + + +def _regular_file_identity(path: Path, root: Path) -> dict[str, object]: + try: + relative = path.relative_to(root) + except ValueError as error: + raise ValueError(f"source closure path escapes its root: {path}") from error + cursor = root + if root.is_symlink() or not root.is_dir(): + raise ValueError(f"source closure root must be a regular directory: {root}") + for part in relative.parts: + cursor /= part + if cursor.is_symlink(): + raise ValueError(f"source closure rejects symlink traversal: {path}") + if not path.is_file(): + raise ValueError(f"source closure requires a regular non-symlink file: {path}") + return {"bytes": path.stat().st_size, "sha256": _sha256(path)} + + +def _canonical_owned_path(value: str) -> str: + path = Path(value) + normalized = path.as_posix() + if ( + path.is_absolute() + or normalized != value + or not path.parts + or ".." in path.parts + or path.parts[0] in {"fbcode", "xplat"} + ): + raise ValueError(f"non-canonical owned source path: {value}") + return normalized + + +def _run_source_control(argv: Sequence[str], label: str) -> str: + try: + result = subprocess.run( + list(argv), + check=False, + capture_output=True, + text=True, + ) + except OSError as error: + raise ValueError(f"cannot inspect {label} checkout: {error}") from error + if result.returncode != 0: + detail = result.stderr.strip() or result.stdout.strip() + raise ValueError(f"cannot inspect {label} checkout: {detail}") + return result.stdout + + +def _checkout_snapshot(root: Path, kind: str) -> dict[str, object]: + if root.is_symlink() or not root.is_dir(): + raise ValueError(f"{kind} checkout root must be a regular directory") + if kind == "fbsource": + head = _run_source_control( + [ + "sl", + "--cwd", + str(root), + "log", + "-r", + ".", + "-T", + "{node}", + "--reason", + "derive Gemma source receipt - sl help log", + ], + kind, + ).strip() + status = _run_source_control( + [ + "sl", + "--cwd", + str(root), + "status", + "--reason", + "verify Gemma source checkout - sl help status", + ], + kind, + ) + elif kind == "oss": + head = _run_source_control( + ["git", "-C", str(root), "rev-parse", "HEAD"], kind + ).strip() + status = _run_source_control( + [ + "git", + "-C", + str(root), + "status", + "--porcelain", + "--untracked-files=all", + ], + kind, + ) + else: + raise ValueError(f"unsupported source checkout kind: {kind}") + if not _is_hex_digest(head, 40): + raise ValueError(f"{kind} checkout did not report a 40-character head") + if status.strip(): + raise ValueError(f"{kind} checkout is not clean") + return {"clean": True, "head": head} + + +def _derive_owned_paths( + fbsource_root: Path, + reviewed_summaries: Sequence[str] | None = None, +) -> list[str]: + summaries = tuple( + _GEMMA_PRODUCTION_DIFF_SUMMARIES + if reviewed_summaries is None + else reviewed_summaries + ) + if not summaries: + raise ValueError("reviewed Gemma production summaries must not be empty") + expected_reverse = tuple(reversed(summaries)) + changed: set[str] = set() + for offset, expected_summary in enumerate(expected_reverse): + revision = "." if offset == 0 else f".~{offset}" + identity = _run_source_control( + [ + "sl", + "--cwd", + str(fbsource_root), + "log", + "-r", + revision, + "-T", + "{node}\\n{desc|firstline}\\n", + "--reason", + "derive Gemma production ownership - sl help log", + ], + "fbsource", + ).splitlines() + if ( + len(identity) != 2 + or not _is_hex_digest(identity[0], 40) + or identity[1] != expected_summary + ): + raise ValueError("fbsource head is not the reviewed Gemma production stack") + paths = _run_source_control( + [ + "sl", + "--cwd", + str(fbsource_root), + "status", + "--change", + identity[0], + "--no-status", + "--root-relative", + "--reason", + "derive Gemma production file union - sl help status", + ], + "fbsource", + ).splitlines() + changed.update(path for path in paths if path) + xplat_prefix = "xplat/executorch/" + fbcode_prefix = "fbcode/executorch/" + unsupported = sorted( + path + for path in changed + if not path.startswith(xplat_prefix) and not path.startswith(fbcode_prefix) + ) + if unsupported: + raise ValueError( + "Gemma production diffs contain files outside the mirrored ExecuTorch roots" + ) + xplat = { + path.removeprefix(xplat_prefix) + for path in changed + if path.startswith(xplat_prefix) + } + fbcode = { + path.removeprefix(fbcode_prefix) + for path in changed + if path.startswith(fbcode_prefix) + } + if not xplat or xplat != fbcode: + raise ValueError( + "Gemma production xplat/fbcode ownership is not byte-mirror complete" + ) + return sorted(_canonical_owned_path(path) for path in xplat) + + +def _source_copy_paths( + fbsource_root: Path, oss_root: Path, logical_path: str +) -> dict[str, Path]: + return { + "fbcode": fbsource_root / "fbcode/executorch" / logical_path, + "oss": oss_root / logical_path, + "xplat": fbsource_root / "xplat/executorch" / logical_path, + } + + +def create_source_manifest( + fbsource_root: Path, + oss_root: Path, +) -> dict[str, object]: + before = { + "fbsource": _checkout_snapshot(fbsource_root, "fbsource"), + "oss": _checkout_snapshot(oss_root, "oss"), + } + paths = _derive_owned_paths(fbsource_root) + files: list[dict[str, object]] = [] + for logical_path in paths: + copy_paths = _source_copy_paths(fbsource_root, oss_root, logical_path) + copies = { + label: { + "path": path.relative_to(root).as_posix(), + **_regular_file_identity(path, root), + } + for label, path, root in ( + ("fbcode", copy_paths["fbcode"], fbsource_root), + ("oss", copy_paths["oss"], oss_root), + ("xplat", copy_paths["xplat"], fbsource_root), + ) + } + identities = {(copy["bytes"], copy["sha256"]) for copy in copies.values()} + if len(identities) != 1: + raise ValueError(f"source mirror/OSS identity mismatch: {logical_path}") + files.append({"copies": copies, "path": logical_path}) + after = { + "fbsource": _checkout_snapshot(fbsource_root, "fbsource"), + "oss": _checkout_snapshot(oss_root, "oss"), + } + if after != before: + raise ValueError("source checkout changed while creating the manifest") + return { + "checkouts": after, + "file_set_sha256": _canonical_set_digest(paths), + "files": files, + "schema_version": _SOURCE_MANIFEST_SCHEMA_VERSION, + } + + +def _validate_checkout_identity(checkout: object, label: str) -> None: + if not isinstance(checkout, dict): + raise ValueError(f"Gemma4 {label} checkout must be an object") + expected = {"clean": True, "head": checkout.get("head")} + if checkout != expected or not _is_hex_digest(checkout.get("head"), 40): + raise ValueError(f"Gemma4 {label} checkout identity is invalid") + + +def _validated_source_manifest_entry(entry: object) -> str: + if not isinstance(entry, dict): + raise ValueError("Gemma4 source manifest entry must be an object") + _require_exact_keys(entry, {"copies", "path"}, "source manifest entry") + path = entry.get("path") + if not isinstance(path, str): + raise ValueError("Gemma4 source manifest path must be a string") + path = _canonical_owned_path(path) + copies = entry.get("copies") + if not isinstance(copies, dict): + raise ValueError("Gemma4 source manifest copies must be an object") + _require_exact_keys(copies, {"fbcode", "oss", "xplat"}, "source copies") + expected_paths = { + "fbcode": f"fbcode/executorch/{path}", + "oss": path, + "xplat": f"xplat/executorch/{path}", + } + identities: set[tuple[object, object]] = set() + for label, expected_path in expected_paths.items(): + copy_identity = copies[label] + if not isinstance(copy_identity, dict): + raise ValueError("Gemma4 source manifest copy must be an object") + _require_exact_keys( + copy_identity, {"bytes", "path", "sha256"}, "source copy identity" + ) + valid = ( + copy_identity.get("path") == expected_path + and isinstance(copy_identity.get("bytes"), int) + and int(copy_identity["bytes"]) >= 0 + and _is_hex_digest(copy_identity.get("sha256"), 64) + ) + if not valid: + raise ValueError("Gemma4 source manifest copy identity is invalid") + identities.add((copy_identity["bytes"], copy_identity["sha256"])) + if len(identities) != 1: + raise ValueError("Gemma4 source manifest mirror/OSS identity mismatch") + return path + + +def validate_source_manifest(manifest: Mapping[str, object]) -> None: + _require_exact_keys( + manifest, + {"checkouts", "file_set_sha256", "files", "schema_version"}, + "Gemma4 source manifest", + ) + if manifest.get("schema_version") != _SOURCE_MANIFEST_SCHEMA_VERSION: + raise ValueError("Gemma4 source manifest schema mismatch") + checkouts = manifest.get("checkouts") + if not isinstance(checkouts, dict): + raise ValueError("Gemma4 source manifest checkouts must be an object") + _require_exact_keys(checkouts, {"fbsource", "oss"}, "source checkouts") + for label in ("fbsource", "oss"): + _validate_checkout_identity(checkouts[label], label) + files = manifest.get("files") + if not isinstance(files, list) or not files: + raise ValueError("Gemma4 source manifest files must be a non-empty list") + paths = [_validated_source_manifest_entry(entry) for entry in files] + if paths != sorted(set(paths)): + raise ValueError("Gemma4 source manifest paths are not sorted and unique") + if manifest.get("file_set_sha256") != _canonical_set_digest(paths): + raise ValueError("Gemma4 source manifest file-set identity mismatch") + + +def _load_wgsl_generator(backend_root: Path) -> Any: + generator_path = backend_root / "scripts/gen_wgsl_headers.py" + _regular_file_identity(generator_path, backend_root) + spec = importlib.util.spec_from_file_location( + "_gemma4_wgsl_generator", generator_path + ) + if spec is None: + raise ValueError("cannot load the WGSL generator") + loader = spec.loader + if loader is None: + raise ValueError("cannot load the WGSL generator") + generator = importlib.util.module_from_spec(spec) + loader.exec_module(generator) + generator.BACKEND_ROOT = backend_root.resolve(strict=True) + return generator + + +def create_wgsl_manifest(backend_root: Path) -> dict[str, object]: + backend_root = backend_root.resolve(strict=True) + if backend_root.parts[-4:] != ("xplat", "executorch", "backends", "webgpu"): + raise ValueError("WGSL backend root is not inside an fbsource xplat checkout") + fbsource_root = backend_root.parents[3] + before = _checkout_snapshot(fbsource_root, "fbsource") + generator = _load_wgsl_generator(backend_root) + shaders = list(generator.discover()) + outputs, orphans = generator.collect_outputs() + if orphans: + raise ValueError("WGSL manifest rejects orphan generated headers") + yaml_paths = sorted((backend_root / "runtime/ops").glob("**/*.yaml")) + expected_yaml = sorted( + shader.with_suffix(".yaml") + for shader in shaders + if shader.with_suffix(".yaml").exists() + ) + if yaml_paths != expected_yaml: + raise ValueError("WGSL manifest rejects orphan YAML specifications") + roles: dict[Path, str] = { + backend_root / "scripts/gen_wgsl_headers.py": "generator", + **{shader: "wgsl" for shader in shaders}, + **{path: "yaml" for path in yaml_paths}, + } + for output, generated in outputs.items(): + identity = _regular_file_identity(output, backend_root) + if output.read_bytes() != generated: + raise ValueError(f"WGSL generated output is stale: {output}") + roles[output] = ( + "global_registry" + if output == generator.registry_path() + else "generated_header" + ) + if identity["bytes"] != len(generated): + raise ValueError(f"WGSL generated output size mismatch: {output}") + files = [ + { + "path": path.relative_to(backend_root).as_posix(), + "role": roles[path], + **_regular_file_identity(path, backend_root), + } + for path in sorted(roles) + ] + after = _checkout_snapshot(fbsource_root, "fbsource") + if after != before: + raise ValueError("fbsource checkout changed while creating the WGSL manifest") + path_roles = [{"path": entry["path"], "role": entry["role"]} for entry in files] + return { + "fbsource_commit": after["head"], + "file_set_sha256": _canonical_set_digest(path_roles), + "files": files, + "orphans": [], + "schema_version": _WGSL_MANIFEST_SCHEMA_VERSION, + } + + +def _validated_wgsl_manifest_entry(entry: object) -> tuple[str, str]: + if not isinstance(entry, dict): + raise ValueError("Gemma4 WGSL manifest entry must be an object") + _require_exact_keys( + entry, {"bytes", "path", "role", "sha256"}, "WGSL file identity" + ) + path = entry.get("path") + role = entry.get("role") + valid_roles = {"generated_header", "generator", "global_registry", "wgsl", "yaml"} + valid = ( + isinstance(path, str) + and _canonical_owned_path(path) == path + and role in valid_roles + and isinstance(entry.get("bytes"), int) + and int(entry["bytes"]) >= 0 + and _is_hex_digest(entry.get("sha256"), 64) + ) + if not valid: + raise ValueError("Gemma4 WGSL manifest file identity is invalid") + assert isinstance(path, str) + return path, str(role) + + +def validate_wgsl_manifest(manifest: Mapping[str, object]) -> None: + _require_exact_keys( + manifest, + { + "fbsource_commit", + "file_set_sha256", + "files", + "orphans", + "schema_version", + }, + "Gemma4 WGSL manifest", + ) + if manifest.get("schema_version") != _WGSL_MANIFEST_SCHEMA_VERSION: + raise ValueError("Gemma4 WGSL manifest schema mismatch") + if manifest.get("orphans") != []: + raise ValueError("Gemma4 WGSL manifest contains orphan outputs") + if not _is_hex_digest(manifest.get("fbsource_commit"), 40): + raise ValueError("Gemma4 WGSL manifest fbsource identity is invalid") + files = manifest.get("files") + if not isinstance(files, list) or not files: + raise ValueError("Gemma4 WGSL manifest files must be a non-empty list") + path_roles: list[dict[str, str]] = [] + role_counts: dict[str, int] = {} + for entry in files: + path, role = _validated_wgsl_manifest_entry(entry) + path_roles.append({"path": path, "role": role}) + role_counts[role] = role_counts.get(role, 0) + 1 + if path_roles != sorted(path_roles, key=lambda item: item["path"]): + raise ValueError("Gemma4 WGSL manifest paths are not sorted") + if len({item["path"] for item in path_roles}) != len(path_roles): + raise ValueError("Gemma4 WGSL manifest contains duplicate paths") + if role_counts.get("generator") != 1 or role_counts.get("global_registry") != 1: + raise ValueError("Gemma4 WGSL manifest requires one generator and registry") + if not role_counts.get("wgsl") or not role_counts.get("generated_header"): + raise ValueError("Gemma4 WGSL manifest has incomplete source/output closure") + if manifest.get("file_set_sha256") != _canonical_set_digest(path_roles): + raise ValueError("Gemma4 WGSL manifest file-set identity mismatch") + + +def create_source_closure_receipt( + fbsource_root: Path, oss_root: Path, backend_root: Path +) -> dict[str, object]: + source_manifest = create_source_manifest(fbsource_root, oss_root) + wgsl_manifest = create_wgsl_manifest(backend_root) + checkouts = source_manifest["checkouts"] + assert isinstance(checkouts, dict) + fbsource = checkouts["fbsource"] + oss = checkouts["oss"] + assert isinstance(fbsource, dict) and isinstance(oss, dict) + if wgsl_manifest["fbsource_commit"] != fbsource["head"]: + raise ValueError("source and WGSL manifests describe different fbsource heads") + return { + "fbsource_commit": fbsource["head"], + "oss_commit": oss["head"], + "schema_version": _SOURCE_CLOSURE_RECEIPT_SCHEMA_VERSION, + "source_current": True, + "source_manifest": copy.deepcopy(source_manifest), + "verification": { + "source_checkout": "verified", + "wgsl_codegen": "verified", + }, + "wgsl_manifest": copy.deepcopy(wgsl_manifest), + } + + +def _validate_source_receipt(root: Path, artifacts: Sequence[object]) -> None: + 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("Gemma4 plain manifest requires one source receipt") + receipt = _load_json(root / source_paths[0]) + _require_exact_keys( + receipt, + { + "fbsource_commit", + "oss_commit", + "schema_version", + "source_current", + "source_manifest", + "verification", + "wgsl_manifest", + }, + "Gemma4 source receipt", + ) + if ( + receipt.get("schema_version") != _SOURCE_CLOSURE_RECEIPT_SCHEMA_VERSION + or receipt.get("source_current") is not True + ): + raise ValueError("Gemma4 source receipt is not source-current") + source_manifest = receipt.get("source_manifest") + wgsl_manifest = receipt.get("wgsl_manifest") + if not isinstance(source_manifest, dict) or not isinstance(wgsl_manifest, dict): + raise ValueError("Gemma4 source receipt lacks semantic manifests") + validate_source_manifest(source_manifest) + validate_wgsl_manifest(wgsl_manifest) + checkouts = source_manifest["checkouts"] + assert isinstance(checkouts, dict) + fbsource = checkouts["fbsource"] + oss = checkouts["oss"] + assert isinstance(fbsource, dict) and isinstance(oss, dict) + if wgsl_manifest.get("fbsource_commit") != fbsource.get("head"): + raise ValueError("Gemma4 source receipt has a mismatched WGSL checkout") + if receipt.get("fbsource_commit") != fbsource.get("head"): + raise ValueError("Gemma4 source receipt has an invalid fbsource commit") + if receipt.get("oss_commit") != oss.get("head"): + raise ValueError("Gemma4 source receipt has an invalid OSS commit") + if receipt.get("verification") != { + "source_checkout": "verified", + "wgsl_codegen": "verified", + }: + raise ValueError("Gemma4 source receipt verification is incomplete") + + +def validate_export_identity(checkpoint_root: Path) -> Mapping[str, object]: + source_config = _source_config_path() + validate_manifest( + source_config.parent, + _single_file_manifest( + source_config.name, SOURCE_CONFIG_BYTES, SOURCE_CONFIG_SHA256 + ), + ) + _validate_architecture(_load_json(source_config), "source config") + + files = 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"]), + ), + ) + _validate_architecture( + _load_json(checkpoint_root / "config.json"), "checkpoint config" + ) + return CHECKPOINT_ACQUISITION + + +def create_plain_manifest( + root: Path, + role_paths: Mapping[str, Path], + ptd_paths: Sequence[Path], +) -> dict[str, object]: + if "pte" not in role_paths or "source" not in role_paths: + raise ValueError("Gemma4 plain manifest requires PTE and source receipt roles") + if len(ptd_paths) != 3: + raise ValueError("Gemma4 plain manifest requires exactly three ordered PTDs") + manifest = create_manifest(root, role_paths, ptd_paths) + artifacts = manifest.get("artifacts") + assert isinstance(artifacts, list) + 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 PTD exceeds the WebGPU external-constant limit") + manifest.update( + { + "acquisition": CHECKPOINT_ACQUISITION, + "export": EXPORT_CONTRACT, + "model": { + "architecture": ARCHITECTURE_FINGERPRINT, + "source_config": { + "path": "config/e2b_config.json", + "sha256": SOURCE_CONFIG_SHA256, + }, + }, + } + ) + _validate_source_receipt(root, artifacts) + return manifest + + +def validate_plain_manifest( + root: Path, + manifest: Mapping[str, object], + require_source_receipt: bool = True, +) -> None: + validate_manifest(root, manifest) + 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: + raise ValueError("Gemma4 architecture identity mismatch") + source_config = model.get("source_config") + if source_config != { + "path": "config/e2b_config.json", + "sha256": SOURCE_CONFIG_SHA256, + }: + raise ValueError("Gemma4 source config identity mismatch") + if manifest.get("export") != EXPORT_CONTRACT: + raise ValueError("Gemma4 WebGPU export contract mismatch") + + artifacts = manifest.get("artifacts") + ptd_order = manifest.get("ptd_order") + assert isinstance(artifacts, list) + assert isinstance(ptd_order, list) + if len(ptd_order) != 3: + raise ValueError("Gemma4 plain manifest requires exactly three ordered PTDs") + 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 PTD exceeds the WebGPU external-constant limit") + roles = { + 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") + if require_source_receipt: + _validate_source_receipt(root, artifacts) + + expected_paths = { + 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") + actual_paths = {entry.name for entry in root.iterdir()} + if actual_paths != expected_paths: + raise ValueError("Gemma4 artifact staging contains missing or extra entries") + + +def _role_paths(values: Sequence[str]) -> dict[str, Path]: + result: dict[str, Path] = {} + for value in values: + role, separator, path = value.partition("=") + if not separator or not role or not path or role in result: + raise ValueError(f"invalid or duplicate ROLE=PATH: {value}") + result[role] = Path(path) + return result + + +def _write_json(path: Path, document: object) -> None: + path.write_text( + json.dumps(document, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + + +def _handle_closure_creation(args: argparse.Namespace) -> bool: + if args.command == "create-source-manifest": + _write_json( + args.output, + create_source_manifest(args.fbsource_root, args.oss_root), + ) + return True + if args.command == "create-wgsl-manifest": + _write_json(args.output, create_wgsl_manifest(args.backend_root)) + return True + if args.command == "create-source-receipt": + _write_json( + args.output, + create_source_closure_receipt( + args.fbsource_root, args.oss_root, args.backend_root + ), + ) + return True + return False + + +def main(argv: Sequence[str] | None = None) -> int: + parser = argparse.ArgumentParser() + subparsers = parser.add_subparsers(dest="command", required=True) + acquisition = subparsers.add_parser("validate-acquisition") + 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=[]) + validate = subparsers.add_parser("validate") + validate.add_argument("--root", type=Path, required=True) + validate.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) + create_source.add_argument("--output", type=Path, required=True) + create_wgsl = subparsers.add_parser("create-wgsl-manifest") + create_wgsl.add_argument("--backend-root", type=Path, required=True) + create_wgsl.add_argument("--output", type=Path, required=True) + create_source_receipt = subparsers.add_parser("create-source-receipt") + create_source_receipt.add_argument("--fbsource-root", type=Path, required=True) + 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) + args = parser.parse_args(argv) + + if args.command == "validate-acquisition": + validate_export_identity(args.checkpoint_root) + 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 + ) + args.output.write_text( + json.dumps(manifest, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + return 0 + + manifest = _load_json(args.manifest) + validate_plain_manifest(args.root, manifest) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/models/gemma4/webgpu_partitioner.py b/examples/models/gemma4/webgpu_partitioner.py new file mode 100644 index 00000000000..8e99733fdfe --- /dev/null +++ b/examples/models/gemma4/webgpu_partitioner.py @@ -0,0 +1,202 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-strict + +"""Gemma 4-specific, instance-scoped WebGPU partitioning.""" + +from functools import lru_cache +from typing import Callable, List, Optional, Tuple + +import executorch.backends.vulkan.patterns as vk_patterns +import torch + +from executorch.backends.vulkan.op_registry import get_op_features, OpFeatures, OpKey +from executorch.backends.vulkan.partitioner.vulkan_partitioner import VulkanPartitioner +from executorch.backends.vulkan.patterns.rope_hf import ( + create_hf_rotary_emb_single_custom_op, + HfRotaryEmbeddingSinglePattern, +) +from executorch.exir import EdgeCompileConfig, ExportedProgram, to_edge +from executorch.exir.backend.partitioner import Partitioner, PartitionResult +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import ExportedProgramPassBase, ExportedProgramPassResult + +from .webgpu_artifact_manifest import WEBGPU_EXTERNAL_CONSTANTS_MAX_DATA_BYTES + +_EXPECTED_GEMMA4_SDPA_COUNT = 35 +_EXPECTED_SINGLE_HF_ROPE_COUNT = 20 + + +def _single_hf_rope_features() -> OpFeatures: + return get_op_features(exir_ops.edge.et_vk.apply_rotary_emb_hf.default) + + +def _extra_op_features() -> dict[OpKey, OpFeatures]: + return { + exir_ops.edge.et_vk.apply_rotary_emb_hf_single.default: ( + _single_hf_rope_features() + ), + exir_ops.edge.et_vk.gemma4_sdpa.default: get_op_features("llama::custom_sdpa"), + } + + +def _webgpu_allowlist() -> list[OpKey]: + return [ + exir_ops.edge.aten.add.Tensor, + exir_ops.edge.et_vk.rms_norm.default, + exir_ops.edge.et_vk.apply_rotary_emb_hf.default, + exir_ops.edge.et_vk.apply_rotary_emb_hf_single.default, + exir_ops.edge.et_vk.gemma4_sdpa.default, + exir_ops.edge.aten.mul.Tensor, + exir_ops.edge.dim_order_ops._clone_dim_order.default, + exir_ops.edge.dim_order_ops._to_dim_order_copy.default, + exir_ops.edge.aten.view_copy.default, + exir_ops.edge.aten.select_copy.int, + exir_ops.edge.aten.sigmoid.default, + exir_ops.edge.aten.gelu.default, + exir_ops.edge.aten.clamp.default, + exir_ops.edge.aten.div.Tensor, + exir_ops.edge.aten.tanh.default, + exir_ops.edge.aten.squeeze_copy.dims, + exir_ops.edge.aten.unsqueeze_copy.default, + exir_ops.edge.aten.slice_copy.Tensor, + exir_ops.edge.aten.permute_copy.default, + exir_ops.edge.aten.cat.default, + exir_ops.edge.aten.argmax.default, + exir_ops.edge.aten._assert_scalar.default, + exir_ops.edge.aten.sym_constrain_range_for_size.default, + exir_ops.edge.et_vk.select_as_symint.default, + ] + + +@lru_cache(maxsize=1) +def _single_hf_rope_patterns() -> List[torch.fx.GraphModule]: + x = torch.randn(1, 1, 4, 32, dtype=torch.float32) + freqs_cos = torch.randn(1, 32, dtype=torch.float32) + freqs_sin = torch.randn(1, 32, dtype=torch.float32) + edge = to_edge( + torch.export.export( + HfRotaryEmbeddingSinglePattern(), + (x, freqs_cos, freqs_sin), + strict=True, + ), + compile_config=EdgeCompileConfig(_check_ir_validity=False), + ) + return [edge.exported_program().graph_module] + + +def _replace_single_hf_rope(exported_program: ExportedProgram) -> None: + graph_module = exported_program.graph_module + vk_patterns.create_replacement_for_pattern( + exported_program, + graph_module, + _single_hf_rope_patterns(), + create_hf_rotary_emb_single_custom_op, + ) + replaced = sum( + node.target == exir_ops.edge.et_vk.apply_rotary_emb_hf_single.default + for node in graph_module.graph.nodes + ) + if replaced != _EXPECTED_SINGLE_HF_ROPE_COUNT: + raise ValueError( + "Gemma4 WebGPU expected " + f"{_EXPECTED_SINGLE_HF_ROPE_COUNT} single-HF-RoPE sites, got {replaced}" + ) + + +def _rank(node: object) -> Optional[int]: + if not isinstance(node, torch.fx.Node): + return None + value = node.meta.get("val") + return value.dim() if isinstance(value, torch.Tensor) else None + + +def _rewrite_gemma4_sdpa(exported_program: ExportedProgram) -> None: + graph_module = exported_program.graph_module + rewritten = 0 + for node in graph_module.graph.nodes: + if node.target != exir_ops.edge.llama.custom_sdpa.default: + continue + if len(node.args) != 8 or node.kwargs: + raise ValueError("Gemma4 custom SDPA must use the exact positional ABI") + query, key, value, _start_pos, mask, dropout, is_causal, scale = node.args + if ( + _rank(query) != 4 + or _rank(key) != 4 + or _rank(value) != 4 + or _rank(mask) != 2 + or dropout != 0.0 + or is_causal is not False + or scale != 1.0 + ): + raise ValueError("Gemma4 custom SDPA call is not WebGPU-compatible") + node.target = exir_ops.edge.et_vk.gemma4_sdpa.default + rewritten += 1 + + if rewritten != _EXPECTED_GEMMA4_SDPA_COUNT: + raise ValueError( + "Gemma4 WebGPU expected " + f"{_EXPECTED_GEMMA4_SDPA_COUNT} SDPA sites, got {rewritten}" + ) + graph_module.recompile() + + +class _Gemma4WebGPURewritePass(ExportedProgramPassBase): + """Plain-Gemma edge rewrites, applied before partitioning. + + `to_edge_transform_and_lower` hands partitioners a deep copy and asserts + the returned graph is identical, so a partitioner cannot rewrite the graph. + The rewrite set is fixed: single-HF-RoPE then Gemma SDPA. MTP extends the + model-owned mechanism separately. + """ + + def call(self, exported_program: ExportedProgram) -> ExportedProgramPassResult: + _replace_single_hf_rope(exported_program) + _rewrite_gemma4_sdpa(exported_program) + return ExportedProgramPassResult(exported_program, True) + + +class Gemma4WebGPUPartitioner(Partitioner): + """Vulkan serialization restricted to Gemma 4 WebGPU capabilities.""" + + def __init__(self, text_quantize: str) -> None: + if "emb8" in text_quantize: + raise ValueError( + "WebGPU cannot delegate emb8; use emb4 (for example, 8da4w+emb4)" + ) + self._inner = VulkanPartitioner( + { + "external_constants_max_data_bytes": ( + WEBGPU_EXTERNAL_CONSTANTS_MAX_DATA_BYTES + ), + "require_dynamic_shapes": True, + "skip_bool_tensors": True, + }, + operator_allowlist=_webgpu_allowlist(), + extra_op_features=_extra_op_features(), + ) + + def ops_to_not_decompose(self, ep: ExportedProgram) -> Tuple[ + List[torch._ops.OpOverload], + Optional[Callable[[torch.fx.Node], bool]], + ]: + return self._inner.ops_to_not_decompose(ep) + + def partition(self, exported_program: ExportedProgram) -> PartitionResult: + # No graph mutation here: to_edge_transform_and_lower asserts the + # partitioner returns an identical graph. The rewrites run in + # Gemma4WebGPURewritePass, before partitioning. + return self._inner.partition(exported_program) + + +def build_webgpu_partitioner(text_quantize: str) -> Gemma4WebGPUPartitioner: + return Gemma4WebGPUPartitioner(text_quantize) + + +def build_webgpu_transform_passes() -> List[ExportedProgramPassBase]: + """Edge transform passes the WebGPU text decoder must run pre-partition.""" + return [_Gemma4WebGPURewritePass()]