diff --git a/examples/models/gemma4/BUCK b/examples/models/gemma4/BUCK index f6b388bc04b..bf44a826628 100644 --- a/examples/models/gemma4/BUCK +++ b/examples/models/gemma4/BUCK @@ -52,6 +52,64 @@ fbcode_target(_kind = runtime.python_binary, define_webgpu_python_targets() +fbcode_target(_kind = runtime.python_library, + name = "mtp_export_lib", + srcs = [ + "eagle_webgpu_round.py", + "export_assistant_webgpu_artifacts.py", + "export_speculative.py", + ], + _is_external_target = True, + base_module = "executorch.examples.models.gemma4", + resources = { + "manifests/gemma4_e2b_mtp_webgpu.json": "manifests/gemma4_e2b_mtp_webgpu.json", + }, + typing = True, + deps = [ + ":quant_utils", + ":text_decoder", + ":webgpu_support", + "//caffe2:torch", + "//executorch/backends/vulkan:custom_ops_lib", + "//executorch/examples/models/llama:source_transformation", + "//executorch/exir:lib", + "//executorch/extension/llm/custom_ops:custom_ops_aot_lib", + "//executorch/extension/llm/custom_ops:custom_ops_aot_py", + "//executorch/extension/llm/export:export_lib", + "//executorch/kernels/quantized:aot_lib", + "//pytorch/ao:torchao", + "fbsource//third-party/pypi/safetensors:safetensors", + "fbsource//third-party/pypi/transformers:transformers", + ], + visibility = ["PUBLIC"], +) + +fbcode_target(_kind = runtime.python_binary, + name = "export_assistant_webgpu_artifacts", + main_function = "executorch.examples.models.gemma4.export_assistant_webgpu_artifacts.main", + preload_deps = [ + "//executorch/extension/llm/custom_ops:custom_ops_aot_lib", + "//executorch/extension/llm/custom_ops:custom_ops_aot_py", + "//executorch/kernels/quantized:aot_lib", + "//pytorch/ao/torchao/csrc/cpu/shared_kernels/embedding_xbit:op_embedding_xbit_aten", + "//pytorch/ao/torchao/csrc/cpu/shared_kernels/linear_8bit_act_xbit_weight:op_linear_8bit_act_xbit_weight_aten", + ], + deps = [":mtp_export_lib"], +) + +fbcode_target(_kind = runtime.python_binary, + name = "export_speculative", + main_function = "executorch.examples.models.gemma4.export_speculative.main", + preload_deps = [ + "//executorch/extension/llm/custom_ops:custom_ops_aot_lib", + "//executorch/extension/llm/custom_ops:custom_ops_aot_py", + "//executorch/kernels/quantized:aot_lib", + "//pytorch/ao/torchao/csrc/cpu/shared_kernels/embedding_xbit:op_embedding_xbit_aten", + "//pytorch/ao/torchao/csrc/cpu/shared_kernels/linear_8bit_act_xbit_weight:op_linear_8bit_act_xbit_weight_aten", + ], + deps = [":mtp_export_lib"], +) + # Text decoder module fbcode_target(_kind = runtime.python_library, name = "text_decoder", diff --git a/examples/models/gemma4/eagle_webgpu_round.py b/examples/models/gemma4/eagle_webgpu_round.py new file mode 100644 index 00000000000..df00c070afc --- /dev/null +++ b/examples/models/gemma4/eagle_webgpu_round.py @@ -0,0 +1,716 @@ +# 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-unsafe + +from __future__ import annotations + +import math +from typing import Any, cast + +import torch + +from executorch.examples.models.gemma4.mtp_qat_contract import ( + OFFICIAL_QAT_CENTROID_TOP_K, + OFFICIAL_QAT_NUM_CENTROIDS, + OFFICIAL_QAT_SELECTED_TOKEN_COUNT, + OFFICIAL_QAT_TOKENS_PER_CENTROID, + validate_qat_token_ordering, +) + + +def _selection_contract( + centroid_top_k: int, selected_token_count: int +) -> dict[str, int]: + if type(centroid_top_k) is not int or type(selected_token_count) is not int: + raise ValueError("assistant selection dimensions must be integers") + if centroid_top_k != OFFICIAL_QAT_CENTROID_TOP_K: + raise ValueError("official QAT assistant requires centroid top-k=32") + if selected_token_count != OFFICIAL_QAT_SELECTED_TOKEN_COUNT: + raise ValueError("official QAT assistant requires 4096 selected tokens") + return { + "centroidTopK": centroid_top_k, + "numCentroids": OFFICIAL_QAT_NUM_CENTROIDS, + "selectedTokenCount": selected_token_count, + "tokensPerCentroid": OFFICIAL_QAT_TOKENS_PER_CENTROID, + } + + +def select_qat_centroids(scores: torch.Tensor) -> torch.Tensor: + if scores.numel() != OFFICIAL_QAT_NUM_CENTROIDS: + raise ValueError("QAT centroid scores must contain 2048 entries") + return torch.topk( + scores.reshape(-1), OFFICIAL_QAT_CENTROID_TOP_K, sorted=True + ).indices + + +def validate_selected_destinations( + token_ordering: torch.Tensor, selected_centroids: torch.Tensor +) -> torch.Tensor: + evidence = validate_qat_token_ordering(token_ordering) + if evidence["permutationExact"] is not True: + raise ValueError("QAT token ordering must be an exact permutation") + if selected_centroids.numel() != OFFICIAL_QAT_CENTROID_TOP_K: + raise ValueError("QAT selection must contain 32 centroids") + if selected_centroids.dtype not in (torch.int32, torch.int64): + raise ValueError("QAT centroid indices must use an integer dtype") + if torch.any(selected_centroids < 0) or torch.any( + selected_centroids >= OFFICIAL_QAT_NUM_CENTROIDS + ): + raise ValueError("QAT centroid index out of range") + if torch.unique(selected_centroids).numel() != OFFICIAL_QAT_CENTROID_TOP_K: + raise ValueError("QAT selected centroids must be distinct") + + logical = token_ordering.reshape( + OFFICIAL_QAT_NUM_CENTROIDS, OFFICIAL_QAT_TOKENS_PER_CENTROID + ) + destinations = logical[selected_centroids.to(torch.int64)].reshape(-1) + if destinations.numel() != OFFICIAL_QAT_SELECTED_TOKEN_COUNT or ( + torch.unique(destinations).numel() != destinations.numel() + ): + raise ValueError("QAT selected destinations must be pairwise distinct") + return destinations + + +class K2LongestPrefixSelector(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.register_buffer( + "zero", torch.zeros((1,), dtype=torch.float32), persistent=False + ) + self.register_buffer( + "one", torch.ones((1,), dtype=torch.float32), persistent=False + ) + + def forward( + self, + round_tokens: torch.Tensor, + greedy_predictions: torch.Tensor, + target_features: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + round_tokens_fp32 = round_tokens.to(torch.float32) + greedy_fp32 = greedy_predictions.to(torch.float32) + target_features_fp32 = target_features.to(torch.float32) + candidates_fp32 = round_tokens_fp32.narrow(1, 1, 2) + + first_matches = (candidates_fp32[:, 0] - greedy_fp32[:, 0]) == 0.0 + second_matches = (candidates_fp32[:, 1] - greedy_fp32[:, 1]) == 0.0 + two = self.one + self.one + accepted_count_fp32 = torch.where( + first_matches, + torch.where(second_matches, two, self.one), + self.zero, + ) + bonus_fp32 = torch.where( + first_matches, + torch.where(second_matches, greedy_fp32[:, 2], greedy_fp32[:, 1]), + greedy_fp32[:, 0], + ) + selected_feature = torch.where( + first_matches, + torch.where( + second_matches, + target_features_fp32[:, 2], + target_features_fp32[:, 1], + ), + target_features_fp32[:, 0], + ) + return ( + accepted_count_fp32.to(torch.long), + bonus_fp32.to(torch.long), + candidates_fp32.to(torch.long), + selected_feature, + ) + + +class Gemma4K2Target(torch.nn.Module): + def __init__(self, text_model: Any) -> None: + super().__init__() + self.text_model: Any = text_model + donors: dict[bool, list[int]] = {False: [], True: []} + for index, layer in enumerate(text_model.self_decoder.layers): + attention = layer.self_attn + if getattr(attention, "is_kv_donor_layer", False): + if attention.kv_cache is None: + raise ValueError("Gemma4 K=2 donor has no KV cache") + donors[bool(attention.is_sliding)].append(index) + if any(len(indices) != 1 for indices in donors.values()): + raise ValueError(f"Gemma4 K=2 donor topology mismatch: {donors}") + if donors != {False: [14], True: [13]}: + raise ValueError(f"Gemma4 K=2 donor layer mismatch: {donors}") + self.full_donor_index = donors[False][0] + self.sliding_donor_index = donors[True][0] + + def donor_views( + self, donor_length: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + length = cast(int, donor_length[0, 0].item()) + torch._check_is_size(length) + torch._check(length >= 2) + full = self.text_model.self_decoder.layers[ + self.full_donor_index + ].self_attn.kv_cache + sliding = self.text_model.self_decoder.layers[ + self.sliding_donor_index + ].self_attn.kv_cache + return ( + full.k_cache.narrow(1, 0, length).transpose(1, 2), + full.v_cache.narrow(1, 0, length).transpose(1, 2), + sliding.k_cache.narrow(1, 0, length).transpose(1, 2), + sliding.v_cache.narrow(1, 0, length).transpose(1, 2), + ) + + def forward( + self, + input_ids: torch.Tensor, + input_pos: torch.Tensor, + is_round: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + mode = cast(int, is_round[0].item()) + torch._check(mode >= 0, lambda: "mode must be 0 or 1") + torch._check_is_size(mode) + torch._check(mode <= 1, lambda: "mode must be 0 or 1") + m = 1 + 2 * mode + length = input_pos.size(0) + torch._check(length >= m, lambda: "round input has too few rows") + torch._check( + length <= 3 + (1 - mode) * (self.text_model.config.max_seq_len - 3), + lambda: "round input must have exactly three rows", + ) + terminal_pos = cast(int, input_pos[-1].item()) + query_start_pos = terminal_pos - m + 1 + torch._check(query_start_pos >= 0, lambda: "positions out of capacity") + torch._check( + terminal_pos < self.text_model.config.max_seq_len, + lambda: "positions out of capacity", + ) + + hidden_states, per_layer_inputs, shared_kv = self.text_model.self_decoder( + input_ids=input_ids, + input_pos=input_pos, + inputs_embeds=None, + ) + hidden_length = hidden_states.size(1) + per_layer_length = per_layer_inputs.size(2) + torch._check( + hidden_length == length, + lambda: "self decoder hidden length must match input length", + ) + torch._check( + per_layer_length == length, + lambda: "self decoder per-layer length must match input length", + ) + hidden_states = torch.narrow(hidden_states, 1, hidden_length - m, m) + per_layer_inputs = torch.narrow(per_layer_inputs, 2, per_layer_length - m, m) + hidden_states = self.text_model.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, + ) + all_features = self.text_model.norm(hidden_states) + all_logits = self.text_model.lm_head(all_features) + features = torch.ops.aten.slice.Tensor( + torch.cat((all_features, all_features, all_features), dim=1), + 1, + 0, + 3, + ) + all_greedy = torch.argmax(all_logits, dim=-1).to(torch.float32) + greedy = torch.ops.aten.slice.Tensor( + torch.cat((all_greedy, all_greedy, all_greedy), dim=1), + 1, + 0, + 3, + ).to(torch.long) + return greedy, features + + +class K2GPUResidentRound(torch.nn.Module): + def __init__( + self, + target: Any, + embed_tokens: Any, + assistant: Any, + hidden_size: int, + max_input_len: int, + max_donor_len: int, + embed_scale: float, + ) -> None: + super().__init__() + if hidden_size <= 0 or max_input_len < 3 or max_donor_len < 2: + raise ValueError("invalid K=2 combined-round dimensions") + if not math.isfinite(embed_scale) or embed_scale <= 0.0: + raise ValueError("invalid Gemma 4 target embedding scale") + self.target: Any = target + self.embed_tokens: Any = embed_tokens + self.assistant: Any = assistant + self.selector = K2LongestPrefixSelector() + self.max_donor_len = max_donor_len + self.embed_scale = embed_scale + self.register_buffer( + "seed_feature", + torch.zeros((1, 1, 1, hidden_size), dtype=torch.float32), + ) + self.register_buffer( + "round_tail", + torch.zeros((1, max_input_len - 3), dtype=torch.float32), + persistent=False, + ) + + def update_seed_feature(self, next_feature: torch.Tensor) -> torch.Tensor: + torch.ops.llama.update_cache.default( + next_feature.unsqueeze(2), self.seed_feature, 0 + ) + return self.seed_feature + + def forward( + self, + input_ids: torch.Tensor, + input_pos: torch.Tensor, + is_round: torch.Tensor, + donor_length: torch.Tensor, + ) -> tuple[torch.Tensor, ...]: + donor_k = cast(int, donor_length[0, 0].item()) + torch._check(donor_k >= 2, lambda: "donor length out of capacity") + torch._check_is_size(donor_k) + torch._check( + donor_k <= self.max_donor_len, + lambda: "donor length out of capacity", + ) + round_condition = is_round.to(torch.float32) == 1.0 + full_k, full_v, sliding_k, sliding_v = self.target.donor_views(donor_length) + first_position_ids = donor_length - 1 + seed_token = input_ids.narrow(1, 0, 1) + seed_embedding = self.embed_tokens(seed_token) * self.embed_scale + first_input = torch.cat((seed_embedding, self.seed_feature.squeeze(2)), dim=-1) + first_logits, first_hidden = self.assistant( + first_input, + first_position_ids, + full_k, + full_v, + sliding_k, + sliding_v, + ) + first_candidate = torch.argmax(first_logits, dim=-1) + + second_embedding = self.embed_tokens(first_candidate) * self.embed_scale + second_input = torch.cat((second_embedding, first_hidden), dim=-1) + second_logits, _ = self.assistant( + second_input, + donor_length, + full_k, + full_v, + sliding_k, + sliding_v, + ) + second_candidate = torch.argmax(second_logits, dim=-1) + round_padded = torch.cat( + ( + seed_token.to(torch.float32), + first_candidate.to(torch.float32), + second_candidate.to(torch.float32), + self.round_tail, + ), + dim=1, + ) + round_tokens = round_padded.narrow(1, 0, input_ids.size(1)) + target_ids = torch.where( + round_condition, round_tokens, input_ids.to(torch.float32) + ).to(torch.long) + target_greedy, target_features = self.target(target_ids, input_pos, is_round) + matches, round_bonus, candidates, selected_feature = self.selector( + round_padded, target_greedy, target_features + ) + + prefill_feature = target_features[:, -1:, :] + round_feature = selected_feature.unsqueeze(1) + next_feature = torch.where(round_condition, round_feature, prefill_feature) + updated_seed_feature = self.update_seed_feature(next_feature) + + matches_fp32 = matches.to(torch.float32) + output_matches = torch.where( + round_condition, matches_fp32, self.selector.zero + ).to(torch.long) + output_bonus = torch.where( + round_condition, + round_bonus.view(1, 1).to(torch.float32), + target_greedy[:, -1:].to(torch.float32), + ).to(torch.long) + state_probe = updated_seed_feature.narrow(3, 0, 1).reshape(1, 1).clone() + return ( + candidates, + target_greedy.to(torch.long), + output_matches, + output_bonus, + state_probe, + ) + + +def _require_tensor_contract( + value: Any, + label: str, + expected_shape: tuple[int | None, ...], + expected_dtype: torch.dtype, +) -> None: + if not isinstance(value, torch.Tensor): + raise ValueError(f"K=2 {label} is not a tensor") + if len(value.shape) != len(expected_shape): + raise ValueError(f"K=2 {label} rank mismatch: {tuple(value.shape)}") + for actual, expected in zip(value.shape, expected_shape): + if expected is not None and actual != expected: + raise ValueError(f"K=2 {label} shape mismatch: {tuple(value.shape)}") + if value.dtype != expected_dtype: + raise ValueError(f"K=2 {label} dtype mismatch: {value.dtype}") + + +def _tensor_meta(value: object) -> torch.Tensor | None: + if not isinstance(value, torch.fx.Node): + return None + tensor = value.meta.get("val") + return tensor if isinstance(tensor, torch.Tensor) else None + + +def _normalize_mutation_target(target: object) -> str: + value = str(target) + if value.endswith("seed_feature"): + return "seed_feature" + marker = "self_decoder.layers." + if marker not in value: + return value + return value[value.index(marker) :] + + +def _expected_mutation_contract(max_donor_len: int) -> list[dict[str, Any]]: + records: list[dict[str, Any]] = [ + { + "logicalTarget": "seed_feature", + "role": "nextFeatureSeed", + "shape": [1, 1, 1, 1536], + "logicalLayout": "BSHD", + "logicalDimOrder": [0, 1, 2, 3], + "vulkanSourceStorage": "BUFFER", + "vulkanDestinationStorage": "TEXTURE_3D", + } + ] + for layer in range(15): + head_dim = 512 if layer in {4, 9, 14} else 256 + for cache_kind in ("k_cache", "v_cache"): + records.append( + { + "logicalTarget": ( + f"self_decoder.layers.{layer}.self_attn.kv_cache." + f"{cache_kind}" + ), + "role": "targetKvCache", + "layer": layer, + "cacheKind": cache_kind, + "shape": [1, max_donor_len, 1, head_dim], + "logicalLayout": "BSHD", + "logicalDimOrder": [0, 1, 2, 3], + "vulkanSourceStorage": "BUFFER", + "vulkanDestinationStorage": "BUFFER", + } + ) + return records + + +def _validate_mutation_contract( + program: torch.export.ExportedProgram, max_donor_len: int +) -> list[dict[str, Any]]: + signature = program.graph_signature + output_specs = list(signature.output_specs) + expected_kinds = ["BUFFER_MUTATION"] * 31 + ["USER_OUTPUT"] * 5 + actual_kinds = [spec.kind.name for spec in output_specs] + if actual_kinds != expected_kinds: + raise ValueError(f"K=2 output-spec order mismatch: {actual_kinds}") + expected = _expected_mutation_contract(max_donor_len) + mutations = output_specs[: len(expected)] + actual_targets = [_normalize_mutation_target(spec.target) for spec in mutations] + expected_targets = [str(record["logicalTarget"]) for record in expected] + if actual_targets != expected_targets: + raise ValueError( + "K=2 mutation target order mismatch: " + f"{actual_targets} != {expected_targets}" + ) + graph_nodes = {node.name: node for node in program.graph.nodes} + for spec, record in zip(mutations, expected): + name = getattr(spec.arg, "name", None) + node = graph_nodes.get(name) + if node is None: + raise ValueError(f"K=2 mutation output is missing: {name}") + _require_tensor_contract( + node.meta.get("val"), + str(record["logicalTarget"]), + tuple(record["shape"]), + torch.float32, + ) + return expected + + +def _donor_view_contract() -> list[dict[str, object]]: + return [ + {"role": "fullK", "layer": 14, "cacheKind": "k_cache", "layout": "BHKD"}, + {"role": "fullV", "layer": 14, "cacheKind": "v_cache", "layout": "BHKD"}, + { + "role": "slidingK", + "layer": 13, + "cacheKind": "k_cache", + "layout": "BHKD", + }, + { + "role": "slidingV", + "layer": 13, + "cacheKind": "v_cache", + "layout": "BHKD", + }, + ] + + +def _validate_seed_alias( + program: torch.export.ExportedProgram, +) -> dict[str, object]: + signature = program.graph_signature + seed_inputs = [ + node + for node in program.graph.nodes + if node.op == "placeholder" + and str(signature.inputs_to_buffers.get(node.name, "")).endswith("seed_feature") + ] + if len(seed_inputs) != 1: + raise ValueError("K=2 seed-feature buffer alias mismatch") + seed = seed_inputs[0] + updates = [ + node + for node in program.graph.nodes + if node.op == "call_function" + and node.target == torch.ops.llama.update_cache.default + and len(node.args) == 3 + and node.args[1] is seed + ] + if len(updates) != 1 or updates[0].args[2] != 0: + raise ValueError("K=2 seed-feature mutation binding mismatch") + source = updates[0].args[0] + _require_tensor_contract( + _tensor_meta(source), + "seed mutation source", + (1, 1, 1, 1536), + torch.float32, + ) + if ( + not isinstance(source, torch.fx.Node) + or source.target != torch.ops.aten.unsqueeze.default + or not source.args + ): + raise ValueError("K=2 seed-feature source must be nextFeature unsqueeze") + _require_tensor_contract( + _tensor_meta(source.args[0]), + "nextFeature", + (1, 1, 1536), + torch.float32, + ) + return { + "logicalSource": "nextFeature[1,1,1536]", + "physicalDestination": "seed_feature[1,1,1,1536]", + "mutation": "llama.update_cache.default", + } + + +def validate_k2_round_abi( # noqa: C901 + program: torch.export.ExportedProgram, + *, + max_input_len: int, + max_donor_len: int, +) -> dict[str, Any]: + signature = program.graph_signature + expected_inputs = ("input_ids", "input_pos", "is_round", "donor_length") + if tuple(signature.user_inputs) != expected_inputs: + raise ValueError( + f"K=2 user-input order mismatch: {tuple(signature.user_inputs)}" + ) + graph_nodes = {node.name: node for node in program.graph.nodes} + input_contracts = ( + ("input_ids", (1, None), torch.int64), + ("input_pos", (None,), torch.int64), + ("is_round", (1,), torch.int64), + ("donor_length", (1, 1), torch.int64), + ) + for name, shape, dtype in input_contracts: + node = graph_nodes.get(name) + if node is None or node.op != "placeholder": + raise ValueError(f"K=2 missing user input: {name}") + _require_tensor_contract(node.meta.get("val"), name, shape, dtype) + input_ids = graph_nodes["input_ids"].meta["val"] + input_pos = graph_nodes["input_pos"].meta["val"] + if str(input_ids.shape[1]) != str(input_pos.shape[0]): + raise ValueError("K=2 input_ids/input_pos dynamic dimensions differ") + + user_outputs = tuple(signature.user_outputs) + if len(user_outputs) != 5: + raise ValueError(f"K=2 user-output count mismatch: {len(user_outputs)}") + output_contracts = ( + ((1, 2), torch.int64, "candidates"), + ((1, 3), torch.int64, "target_greedy"), + ((1,), torch.int64, "output_matches"), + ((1, 1), torch.int64, "output_bonus"), + ((1, 1), torch.float32, "state_probe"), + ) + for name, (shape, dtype, label) in zip(user_outputs, output_contracts): + node = graph_nodes.get(name) + if node is None: + raise ValueError(f"K=2 missing user output: {name}") + _require_tensor_contract(node.meta.get("val"), label, shape, dtype) + + mutation_contract = _validate_mutation_contract(program, max_donor_len) + + operator_counts: dict[str, int] = {} + for node in program.graph.nodes: + if node.op == "call_function": + target = str(node.target) + operator_counts[target] = operator_counts.get(target, 0) + 1 + expected_operator_counts = { + "aten.argmax.default": 3, + "aten.scatter.src": 2, + "aten.topk.default": 2, + "llama.custom_sdpa.default": 43, + "llama.update_cache.default": 31, + } + for target, expected in expected_operator_counts.items(): + if operator_counts.get(target, 0) != expected: + raise ValueError( + f"K=2 {target} count mismatch: " + f"{operator_counts.get(target, 0)} != {expected}" + ) + + range_bounds: set[tuple[int, int]] = set() + for value in program.range_constraints.values(): + try: + bounds = (int(value.lower), int(value.upper)) + except (TypeError, ValueError, OverflowError): + continue + range_bounds.add(bounds) + for expected_range in ((1, max_input_len), (2, max_donor_len)): + if expected_range not in range_bounds: + raise ValueError(f"K=2 missing dynamic range: {expected_range}") + return { + "bufferMutationCount": len(mutation_contract), + "donorViewOrder": _donor_view_contract(), + "inputOrder": [contract[0] for contract in input_contracts], + "mutationOrder": mutation_contract, + "operatorCounts": expected_operator_counts, + "outputOrder": [contract[2] for contract in output_contracts], + "seedMutationCount": 1, + "stateAlias": _validate_seed_alias(program), + } + + +def export_k2_round_program( + module: K2GPUResidentRound, max_input_len: int = 64 +) -> torch.export.ExportedProgram: + if max_input_len < 3 or module.round_tail.numel() + 3 < max_input_len: + raise ValueError("invalid K=2 round input bound") + seq_len = torch.export.Dim("seq_len", min=1, max=max_input_len) + return torch.export.export( + module, + ( + torch.ones((1, 3), dtype=torch.long), + torch.arange(3, dtype=torch.long), + torch.tensor([0], dtype=torch.long), + torch.tensor([[2]], dtype=torch.long), + ), + dynamic_shapes={ + "input_ids": {1: seq_len}, + "input_pos": {0: seq_len}, + "is_round": None, + "donor_length": None, + }, + strict=False, + ) + + +def _rewrite_negative_select_as_symint( + program: torch.export.ExportedProgram, +) -> int: + 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("K=2 negative select_as_symint source must be an input") + source_dtype = getattr(source.meta.get("val"), "dtype", None) + if source_dtype not in {torch.int32, torch.int64}: + raise ValueError("K=2 negative select_as_symint requires integer 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 compose_k2_round_program( + text_model: Any, + embed_tokens: Any, + assistant: Any, + hidden_size: int, + max_seq_len: int, + embed_scale: float, + max_input_len: int = 64, + max_donor_len: int | None = None, +) -> torch.export.ExportedProgram: + donor_bound = max_seq_len if max_donor_len is None else max_donor_len + module = K2GPUResidentRound( + target=Gemma4K2Target(text_model), + embed_tokens=embed_tokens, + assistant=assistant, + hidden_size=hidden_size, + max_input_len=max_input_len, + max_donor_len=donor_bound, + embed_scale=embed_scale, + ).eval() + program = export_k2_round_program(module, max_input_len) + replacements = _rewrite_negative_select_as_symint(program) + if replacements != 2: + raise ValueError( + "K=2 composition requires exactly two select_as_symint rewrites, " + f"found {replacements}" + ) + program = program.run_decompositions({}) + abi = validate_k2_round_abi( + program, + max_input_len=max_input_len, + max_donor_len=donor_bound, + ) + program.graph_module.meta["gemma4K2Abi"] = abi + return program diff --git a/examples/models/gemma4/export_assistant_webgpu_artifacts.py b/examples/models/gemma4/export_assistant_webgpu_artifacts.py new file mode 100644 index 00000000000..eb89cc23b93 --- /dev/null +++ b/examples/models/gemma4/export_assistant_webgpu_artifacts.py @@ -0,0 +1,747 @@ +#!/usr/bin/env python3 +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-unsafe + +from __future__ import annotations + +import argparse +import hashlib +import math +import tempfile +import types +from pathlib import Path +from typing import Any, cast + +import executorch.backends.vulkan.custom_ops_lib # noqa: F401 +import torch + +from executorch.examples.models.gemma4.eagle_webgpu_round import ( + OFFICIAL_QAT_CENTROID_TOP_K, + OFFICIAL_QAT_NUM_CENTROIDS, + OFFICIAL_QAT_SELECTED_TOKEN_COUNT, + OFFICIAL_QAT_TOKENS_PER_CENTROID, + validate_qat_token_ordering, +) +from executorch.examples.models.gemma4.webgpu_artifact_manifest import ( + validate_assistant_export_identity, +) +from safetensors import safe_open + +QAT_VALIDATION_DONOR_SEQUENCE = (2, 16, 511, 512, 513, 514, 1024, 8960, 2) + + +def _tensor_sha256(tensor: torch.Tensor) -> str: + payload = tensor.detach().cpu().contiguous().numpy().tobytes() + return hashlib.sha256(payload).hexdigest() + + +def _load_raw_token_ordering( + checkpoint: Path, +) -> tuple[torch.Tensor, dict[str, Any]]: + model_path = checkpoint / "model.safetensors" + with safe_open(str(model_path), framework="pt", device="cpu") as tensors: + matches = [ + name + for name in tensors.keys() + if name.endswith("masked_embedding.token_ordering") + ] + if len(matches) != 1: + raise ValueError( + f"assistant safetensors token-ordering key mismatch: {matches}" + ) + ordering = tensors.get_tensor(matches[0]) + if ordering.dtype != torch.int64 or tuple(ordering.shape) != (262144,): + raise ValueError( + "assistant raw token ordering must be int64 with shape [262144]" + ) + evidence = validate_qat_token_ordering(ordering) + if evidence["rawShape"] != [262144] or evidence["shape"] != [2048, 128]: + raise ValueError("assistant raw token-ordering shape proof mismatch") + return ordering, evidence + + +def validate_qat_centroid_scores(scores: torch.Tensor) -> dict[str, Any]: + if tuple(scores.shape) != (1, 1, OFFICIAL_QAT_NUM_CENTROIDS): + raise ValueError("assistant centroid scores must have shape [1, 1, 2048]") + if scores.dtype != torch.float32 or not bool(torch.isfinite(scores).all()): + raise ValueError("assistant centroid scores must be finite fp32 values") + + top33_values, top33_indices = torch.topk( + scores, + k=OFFICIAL_QAT_CENTROID_TOP_K + 1, + dim=-1, + sorted=True, + ) + top32_values = top33_values[..., :OFFICIAL_QAT_CENTROID_TOP_K] + top32_indices = top33_indices[..., :OFFICIAL_QAT_CENTROID_TOP_K] + stable_values, stable_indices = torch.sort( + scores, dim=-1, descending=True, stable=True + ) + stable_reference_exact = bool( + torch.equal( + top33_values, + stable_values[..., : OFFICIAL_QAT_CENTROID_TOP_K + 1], + ) + and torch.equal( + top33_indices, + stable_indices[..., : OFFICIAL_QAT_CENTROID_TOP_K + 1], + ) + ) + pairwise_distinct = bool(torch.all(top32_values[..., :-1] > top32_values[..., 1:])) + boundary_gap = float( + top33_values[..., OFFICIAL_QAT_CENTROID_TOP_K - 1] + - top33_values[..., OFFICIAL_QAT_CENTROID_TOP_K] + ) + if not pairwise_distinct: + raise ValueError("assistant top-32 values must be pairwise distinct") + if not math.isfinite(boundary_gap) or boundary_gap <= 0.0: + raise ValueError("assistant 32/33 boundary gap must be positive") + if not stable_reference_exact: + raise ValueError("assistant top-33 order differs from stable descending order") + return { + "allFinite": True, + "boundaryGap": boundary_gap, + "indicesSha256": _tensor_sha256(top32_indices), + "stableReferenceExact": True, + "top32PairwiseDistinct": True, + "top33IndicesSha256": _tensor_sha256(top33_indices), + "top33ValuesSha256": _tensor_sha256(top33_values), + "valuesSha256": _tensor_sha256(top32_values), + } + + +def validate_assistant_checkpoint( + checkpoint: Path, +) -> dict[str, Any]: + checkpoint = checkpoint.resolve(strict=True) + if not checkpoint.is_dir(): + raise ValueError("assistant checkpoint must be a directory") + return dict(validate_assistant_export_identity(checkpoint)) + + +def _quantized_embedding(weight: torch.Tensor, bits: int) -> torch.nn.Module: + from executorch.examples.models.llama.source_transformation.quantize import ( + EmbeddingQuantHandler, + ) + + holder: Any = torch.nn.Module() + holder.embed = torch.nn.Embedding(weight.shape[0], weight.shape[1]) + holder.embed.weight = torch.nn.Parameter(weight) + quantized: Any = EmbeddingQuantHandler( + holder, bitwidth=bits, group_size=None, packed=bits == 4 + ).quantized_model() + return quantized.embed + + +def _masked_embedding_static_output( + self: Any, + hidden_states: torch.Tensor, + _lm_head_weight: torch.Tensor, +) -> torch.Tensor: + batch, seq_len = hidden_states.shape[:2] + centroid_logits = self.centroids(hidden_states) + _, top_k_indices = torch.topk( + centroid_logits, k=OFFICIAL_QAT_CENTROID_TOP_K, dim=-1 + ) + selected_canonical = torch.nn.functional.embedding( + top_k_indices, self._webgpu_token_ordering + ).to(torch.long) + selected_flat = selected_canonical.reshape(-1) + selected_embeddings = self._lm_embed(selected_flat).view( + batch, + seq_len, + OFFICIAL_QAT_SELECTED_TOKEN_COUNT, + hidden_states.shape[-1], + ) + selected_logits = ( + hidden_states.unsqueeze(-2) @ selected_embeddings.transpose(-1, -2) + ).squeeze(-2) + return self._webgpu_output_template.scatter( + dim=-1, + index=selected_canonical.view(batch, seq_len, -1), + src=selected_logits, + ) + + +def adapt_masked_embedding_for_webgpu( + masked_embedding: Any, +) -> Any: + if not hasattr(masked_embedding, "_lm_embed"): + raise ValueError("assistant WebGPU export requires a quantized LM head") + evidence = validate_qat_token_ordering(masked_embedding.token_ordering) + if evidence["shape"] != [ + OFFICIAL_QAT_NUM_CENTROIDS, + OFFICIAL_QAT_TOKENS_PER_CENTROID, + ]: + raise ValueError("assistant token-ordering logical shape mismatch") + dtype = masked_embedding.centroids.weight.dtype + device = masked_embedding.centroids.weight.device + ordering = masked_embedding.token_ordering.detach().to(torch.long).reshape(-1) + masked_embedding.register_buffer( + "_webgpu_token_ordering", + ordering.to(dtype=torch.float32, device=device).view( + OFFICIAL_QAT_NUM_CENTROIDS, + OFFICIAL_QAT_TOKENS_PER_CENTROID, + ), + persistent=False, + ) + masked_embedding.register_buffer( + "_webgpu_output_template", + torch.full( + (1, 1, OFFICIAL_QAT_NUM_CENTROIDS * OFFICIAL_QAT_TOKENS_PER_CENTROID), + torch.finfo(dtype).min, + dtype=dtype, + device=device, + ), + persistent=False, + ) + masked_embedding.forward = types.MethodType( + _masked_embedding_static_output, masked_embedding + ) + return masked_embedding + + +class StaticAssistantMasks(torch.nn.Module): + def __init__(self, max_seq_len: int, sliding_window: int = 512) -> None: + super().__init__() + if max_seq_len < 2 or sliding_window <= 0: + raise ValueError("invalid assistant attention-mask capacity") + self.max_seq_len = max_seq_len + self.register_buffer("full_mask", torch.zeros(max_seq_len), persistent=False) + sliding_mask = torch.full((max_seq_len,), torch.finfo(torch.float32).min) + sliding_mask[-(sliding_window + 1) :] = 0 + self.register_buffer("sliding_mask", sliding_mask, persistent=False) + + def forward( + self, full_k: torch.Tensor, sliding_k: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + full_length = full_k.shape[2] + sliding_length = sliding_k.shape[2] + full_mask = self.full_mask.narrow(0, 0, full_length).unsqueeze(0) + sliding_mask = self.sliding_mask.narrow( + 0, self.max_seq_len - sliding_length, sliding_length + ).unsqueeze(0) + return full_mask, sliding_mask + + +class _StaticAssistantQueryRopeLayer(torch.nn.Module): + def __init__(self, freqs_cos: torch.Tensor, freqs_sin: torch.Tensor) -> None: + super().__init__() + self.register_buffer("freqs_cos", freqs_cos, persistent=False) + self.register_buffer("freqs_sin", freqs_sin, persistent=False) + + def forward(self, query: torch.Tensor, position_ids: torch.Tensor) -> torch.Tensor: + start_pos = cast(int, position_ids[0, 0].item()) + torch._check_is_size(start_pos) + torch._check(start_pos >= 0) + torch._check(start_pos + query.shape[1] <= self.freqs_cos.shape[0]) + return torch.ops.et_vk.apply_rotary_emb_hf_single.default( + query, self.freqs_cos, self.freqs_sin, start_pos + ) + + +class StaticAssistantQueryRope(torch.nn.Module): + def __init__(self, source: Any, max_seq_len: int) -> None: + super().__init__() + source_buffer = next(source.buffers()) + positions = torch.arange(max_seq_len, device=source_buffer.device).unsqueeze(0) + probe = torch.empty(1, 1, 1, dtype=torch.float32, device=source_buffer.device) + with torch.no_grad(): + sliding_cos, sliding_sin = source(probe, positions, "sliding_attention") + full_cos, full_sin = source(probe, positions, "full_attention") + self.sliding_attention = _StaticAssistantQueryRopeLayer( + sliding_cos.squeeze(0).contiguous(), + sliding_sin.squeeze(0).contiguous(), + ) + self.full_attention = _StaticAssistantQueryRopeLayer( + full_cos.squeeze(0).contiguous(), + full_sin.squeeze(0).contiguous(), + ) + + +class StaticAssistantSharedKVAttention(torch.nn.Module): + def __init__( + self, + source: Any, + query_rope: _StaticAssistantQueryRopeLayer, + ) -> None: + super().__init__() + if not source.is_kv_shared_layer: + raise ValueError("assistant attention requires shared target KV") + self.layer_type = source.layer_type + self.head_dim = source.head_dim + self.num_attention_heads = ( + source.num_attention_heads + if hasattr(source, "num_attention_heads") + else source.config.num_attention_heads + ) + self.q_proj: Any = source.q_proj + self.q_norm: Any = source.q_norm + self.o_proj: Any = source.o_proj + self.query_rope = query_rope + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: Any, + attention_mask: torch.Tensor, + shared_kv_states: dict[str, tuple[torch.Tensor, torch.Tensor]], + position_ids: torch.Tensor, + **kwargs: Any, + ) -> tuple[torch.Tensor, None]: + del position_embeddings, kwargs + input_shape = hidden_states.shape[:-1] + query = self.q_proj(hidden_states).view( + *input_shape, self.num_attention_heads, self.head_dim + ) + query = self.q_norm(query) + query = self.query_rope(query, position_ids) + key, value = shared_kv_states[self.layer_type] + output = torch.ops.llama.custom_sdpa.default( + query, + key.transpose(1, 2), + value.transpose(1, 2), + 0, + attention_mask, + 0.0, + False, + 1.0, + ) + return self.o_proj(output.reshape(*input_shape, -1)), None + + +class _UnusedAssistantRotaryEmbedding(torch.nn.Module): + def forward( + self, + x: torch.Tensor, + position_ids: torch.Tensor, + layer_type: str | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + del position_ids, layer_type + return x, x + + +def _static_assistant_attention_masks( + self: Any, + inputs_embeds: torch.Tensor, + attention_mask: torch.Tensor | None, + shared_kv_states: dict[str, tuple[torch.Tensor, torch.Tensor]], +) -> dict[str, torch.Tensor]: + del inputs_embeds + if attention_mask is not None: + raise ValueError("assistant WebGPU export requires attention_mask=None") + full_mask, sliding_mask = self._webgpu_static_masks( + shared_kv_states["full_attention"][0], + shared_kv_states["sliding_attention"][0], + ) + return { + "full_attention": full_mask, + "sliding_attention": sliding_mask, + } + + +def adapt_assistant_model_for_webgpu(assistant: Any, max_seq_len: int) -> Any: + layers = list(assistant.model.layers) + layer_types = [layer.self_attn.layer_type for layer in layers] + expected_layer_types = [ + "sliding_attention", + "sliding_attention", + "sliding_attention", + "full_attention", + ] + if layer_types != expected_layer_types: + raise ValueError(f"unexpected assistant layer types: {layer_types}") + query_rope = StaticAssistantQueryRope( + assistant.model.rotary_emb, max_seq_len=max_seq_len + ) + for layer in layers: + source_rope = getattr(query_rope, layer.self_attn.layer_type) + layer.self_attn = StaticAssistantSharedKVAttention( + layer.self_attn, + _StaticAssistantQueryRopeLayer( + source_rope.freqs_cos, source_rope.freqs_sin + ), + ) + assistant.model.rotary_emb = _UnusedAssistantRotaryEmbedding() + assistant._webgpu_static_masks = StaticAssistantMasks(max_seq_len) + assistant.create_attention_masks = types.MethodType( + _static_assistant_attention_masks, assistant + ) + adapt_masked_embedding_for_webgpu(assistant.masked_embedding) + return assistant + + +class UnfoldedAssistant(torch.nn.Module): + def __init__(self, assistant: Any) -> None: + super().__init__() + self.assistant: Any = assistant + self._webgpu_qat_selection_evidence: dict[str, Any] = {} + + def forward( + self, + inputs_embeds: torch.Tensor, + position_ids: torch.Tensor, + full_k: torch.Tensor, + full_v: torch.Tensor, + sliding_k: torch.Tensor, + sliding_v: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + output = self.assistant( + inputs_embeds=inputs_embeds, + attention_mask=None, + position_ids=position_ids, + shared_kv_states={ + "full_attention": (full_k, full_v), + "sliding_attention": (sliding_k, sliding_v), + }, + use_cache=False, + ) + return output.logits, output.last_hidden_state + + +def _qat_validation_inputs(config: Any, donor_length: int) -> tuple[torch.Tensor, ...]: + text_config = config.get_text_config() + generator = torch.Generator().manual_seed(0xE4A61000 + donor_length) + return ( + torch.randn( + 1, + 1, + 2 * config.backbone_hidden_size, + generator=generator, + ), + torch.tensor([[donor_length - 1]], dtype=torch.long), + torch.randn( + 1, + 1, + donor_length, + text_config.global_head_dim, + generator=generator, + ), + torch.randn( + 1, + 1, + donor_length, + text_config.global_head_dim, + generator=generator, + ), + torch.randn( + 1, + 1, + donor_length, + text_config.head_dim, + generator=generator, + ), + torch.randn( + 1, + 1, + donor_length, + text_config.head_dim, + generator=generator, + ), + ) + + +def _qat_validation_cases( + config: Any, max_donor_len: int +) -> tuple[tuple[int, ...], list[tuple[torch.Tensor, ...]]]: + donor_sequence = tuple( + donor_length + for donor_length in QAT_VALIDATION_DONOR_SEQUENCE + if donor_length <= max_donor_len + ) + if not donor_sequence or donor_sequence[0] != 2 or donor_sequence[-1] != 2: + raise ValueError("assistant validation requires replay at donor length 2") + return donor_sequence, [ + _qat_validation_inputs(config, donor_length) for donor_length in donor_sequence + ] + + +def _capture_reference_outputs( + wrapper: UnfoldedAssistant, + inputs: list[tuple[torch.Tensor, ...]], +) -> list[tuple[torch.Tensor, torch.Tensor]]: + with torch.no_grad(): + return [ + tuple(output.detach().clone() for output in wrapper(*case)) + for case in inputs + ] + + +def validate_qat_selection_contract( # noqa: C901 + wrapper: UnfoldedAssistant, + max_donor_len: int, + *, + validation_inputs: list[tuple[torch.Tensor, ...]] | None = None, + reference_outputs: list[tuple[torch.Tensor, torch.Tensor]] | None = None, + token_ordering_evidence: dict[str, Any] | None = None, +) -> dict[str, Any]: + donor_sequence, default_inputs = _qat_validation_cases( + wrapper.assistant.config, max_donor_len + ) + inputs = default_inputs if validation_inputs is None else validation_inputs + if len(inputs) != len(donor_sequence): + raise ValueError("assistant validation input count mismatch") + if reference_outputs is not None and len(reference_outputs) != len(inputs): + raise ValueError("assistant eager-reference count mismatch") + + masked_embedding: Any = wrapper.assistant.masked_embedding + cases = [] + with torch.no_grad(): + for case_index, (donor_length, case_inputs) in enumerate( + zip(donor_sequence, inputs) + ): + head_inputs: list[torch.Tensor] = [] + hook = masked_embedding.register_forward_pre_hook( + lambda _module, args, destination=head_inputs: destination.append( + args[0].detach() + ) + ) + try: + actual = tuple(wrapper(*case_inputs)) + finally: + hook.remove() + if len(head_inputs) != 1: + raise ValueError( + "assistant masked head must run exactly once per validation case" + ) + output_evidence = [] + greedy_token_exact = True + if reference_outputs is not None: + expected = reference_outputs[case_index] + for name, expected_tensor, actual_tensor in zip( + ("logits", "last_hidden_state"), expected, actual + ): + difference = (actual_tensor - expected_tensor).abs() + close = torch.allclose( + actual_tensor, expected_tensor, rtol=1e-3, atol=1e-4 + ) + if not close: + raise ValueError( + f"assistant adapter {name} mismatch at K={donor_length}: " + f"max_abs={float(difference.max())}" + ) + output_evidence.append( + { + "actualSha256": _tensor_sha256(actual_tensor), + "bitExact": torch.equal(actual_tensor, expected_tensor), + "close": True, + "maxAbsError": float(difference.max()), + "name": name, + "referenceSha256": _tensor_sha256(expected_tensor), + "shape": list(actual_tensor.shape), + } + ) + greedy_token_exact = torch.equal( + actual[0].argmax(dim=-1), expected[0].argmax(dim=-1) + ) + if not greedy_token_exact: + raise ValueError( + f"assistant adapter greedy token mismatch at K={donor_length}" + ) + cases.append( + { + "caseIndex": case_index, + "donorLength": donor_length, + "greedyTokenExact": greedy_token_exact, + "inputSha256": [_tensor_sha256(value) for value in case_inputs], + "outputs": output_evidence, + "topk": validate_qat_centroid_scores( + masked_embedding.centroids(head_inputs[0]) + ), + } + ) + + if cases[0]["inputSha256"] != cases[-1]["inputSha256"]: + raise ValueError("assistant validation replay input mismatch") + if cases[0]["topk"] != cases[-1]["topk"]: + raise ValueError("assistant validation replay top-k mismatch") + if cases[0]["outputs"] != cases[-1]["outputs"]: + raise ValueError("assistant validation replay output mismatch") + evidence: dict[str, Any] = { + "cases": cases, + "donorSequence": list(donor_sequence), + "selectionContract": { + "centroidTopK": OFFICIAL_QAT_CENTROID_TOP_K, + "numCentroids": OFFICIAL_QAT_NUM_CENTROIDS, + "selectedTokenCount": OFFICIAL_QAT_SELECTED_TOKEN_COUNT, + "tokensPerCentroid": OFFICIAL_QAT_TOKENS_PER_CENTROID, + }, + } + if reference_outputs is not None: + evidence["eagerEquivalence"] = { + "allClose": True, + "atol": 1e-4, + "rtol": 1e-3, + } + if token_ordering_evidence is not None: + evidence["tokenOrdering"] = token_ordering_evidence + return evidence + + +def load_qat_assistant( + checkpoint: Path, + *, + max_donor_len: int = 8960, + lm_head_bits: int = 4, + quantize_backbone: str = "8da4w", +) -> UnfoldedAssistant: + if max_donor_len < 2: + raise ValueError("assistant donor capacity must be at least 2") + if lm_head_bits != 4: + raise ValueError("official QAT assistant requires a 4-bit LM head") + if quantize_backbone != "8da4w": + raise ValueError("official QAT assistant requires 8da4w backbone quantization") + validate_assistant_checkpoint(checkpoint) + raw_token_ordering, raw_token_ordering_evidence = _load_raw_token_ordering( + checkpoint + ) + + import importlib + + transformers_module: Any = importlib.import_module("transformers") + assistant_type: Any = transformers_module.Gemma4AssistantForCausalLM + assistant: Any = assistant_type.from_pretrained( + str(checkpoint.resolve()), + torch_dtype=torch.float32, + trust_remote_code=False, + ).eval() + masked_embedding = assistant.masked_embedding + if masked_embedding is None: + raise ValueError("QAT assistant is missing its masked embedding") + if ( + masked_embedding.num_centroids != OFFICIAL_QAT_NUM_CENTROIDS + or masked_embedding.vocab_size_per_centroid != OFFICIAL_QAT_TOKENS_PER_CENTROID + or masked_embedding.vocab_size + != OFFICIAL_QAT_NUM_CENTROIDS * OFFICIAL_QAT_TOKENS_PER_CENTROID + ): + raise ValueError("QAT assistant selection dimensions mismatch") + masked_embedding.centroid_intermediate_top_k = OFFICIAL_QAT_CENTROID_TOP_K + loaded_token_ordering = masked_embedding.token_ordering.detach().cpu().contiguous() + if loaded_token_ordering.dtype != torch.int64 or tuple( + loaded_token_ordering.shape + ) != (262144,): + raise ValueError( + "assistant loaded token ordering must be int64 with shape [262144]" + ) + loaded_token_ordering_evidence = validate_qat_token_ordering(loaded_token_ordering) + raw_sha256 = _tensor_sha256(raw_token_ordering) + loaded_sha256 = _tensor_sha256(loaded_token_ordering) + if ( + raw_sha256 != loaded_sha256 + or raw_sha256 != raw_token_ordering_evidence["sha256"] + or loaded_sha256 != loaded_token_ordering_evidence["sha256"] + or not torch.equal(raw_token_ordering, loaded_token_ordering) + ): + raise ValueError("assistant raw/loaded token ordering differs") + token_ordering_evidence = dict(loaded_token_ordering_evidence) + token_ordering_evidence.update( + { + "loaded": loaded_token_ordering_evidence, + "raw": raw_token_ordering_evidence, + "rawLoadedByteExact": True, + "rawSha256": raw_sha256, + } + ) + masked_embedding._lm_embed = _quantized_embedding( + assistant.lm_head.weight.detach(), lm_head_bits + ) + + from executorch.extension.llm.export.quantize import quantize_model_ + + quantize_model_( + assistant.model, + qlinear_config=quantize_backbone, + qlinear_group_size=128, + skip_incompatible_shapes=True, + ) + _, validation_inputs = _qat_validation_cases(assistant.config, max_donor_len) + reference_outputs = _capture_reference_outputs( + UnfoldedAssistant(assistant).eval(), validation_inputs + ) + adapt_assistant_model_for_webgpu(assistant, max_seq_len=max_donor_len) + wrapper = UnfoldedAssistant(assistant).eval() + wrapper._webgpu_qat_selection_evidence = validate_qat_selection_contract( + wrapper, + max_donor_len, + validation_inputs=validation_inputs, + reference_outputs=reference_outputs, + token_ordering_evidence=token_ordering_evidence, + ) + return wrapper + + +def export_assistant_program( + checkpoint: Path, + *, + max_donor_len: int = 8960, + lm_head_bits: int = 4, + quantize_backbone: str = "8da4w", +) -> torch.export.ExportedProgram: + if max_donor_len < 2: + raise ValueError("assistant donor capacity must be at least 2") + wrapper = load_qat_assistant( + checkpoint, + max_donor_len=max_donor_len, + lm_head_bits=lm_head_bits, + quantize_backbone=quantize_backbone, + ) + config: Any = wrapper.assistant.config + text_config: Any = config.get_text_config() + donor_len = 2 + example = ( + torch.randn(1, 1, 2 * config.backbone_hidden_size), + torch.tensor([[donor_len - 1]], dtype=torch.long), + torch.randn(1, 1, donor_len, text_config.global_head_dim), + torch.randn(1, 1, donor_len, text_config.global_head_dim), + torch.randn(1, 1, donor_len, text_config.head_dim), + torch.randn(1, 1, donor_len, text_config.head_dim), + ) + dynamic_donor = torch.export.Dim("assistant_donor_len", min=2, max=max_donor_len) + with torch.no_grad(): + return torch.export.export( + wrapper, + example, + dynamic_shapes={ + "inputs_embeds": None, + "position_ids": None, + "full_k": {2: dynamic_donor}, + "full_v": {2: dynamic_donor}, + "sliding_k": {2: dynamic_donor}, + "sliding_v": {2: dynamic_donor}, + }, + strict=False, + ) + + +def main() -> int: + parser = argparse.ArgumentParser( + description="Validate and export the Gemma 4 QAT assistant graph" + ) + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--max-donor-len", type=int, default=8960) + args = parser.parse_args() + if args.output.suffix != ".pt2": + raise ValueError("assistant graph output path must end in .pt2") + if args.output.exists(): + raise ValueError(f"refusing to overwrite existing artifact: {args.output}") + program = export_assistant_program( + args.checkpoint, max_donor_len=args.max_donor_len + ) + args.output.parent.mkdir(parents=True, exist_ok=True) + with tempfile.TemporaryDirectory( + prefix=f".{args.output.stem}.", dir=args.output.parent + ) as staging_directory: + staged_output = Path(staging_directory) / args.output.name + torch.export.save(program, staged_output) + staged_output.replace(args.output) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/models/gemma4/export_speculative.py b/examples/models/gemma4/export_speculative.py new file mode 100644 index 00000000000..7a365b7cec5 --- /dev/null +++ b/examples/models/gemma4/export_speculative.py @@ -0,0 +1,396 @@ +#!/usr/bin/env python3 +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +# pyre-unsafe + +from __future__ import annotations + +import argparse +import json +import operator +import shutil +import tempfile +from pathlib import Path +from typing import Any + +import torch + +from executorch.examples.models.gemma4.eagle_webgpu_round import ( + compose_k2_round_program, +) +from executorch.examples.models.gemma4.export_assistant_webgpu_artifacts import ( + load_qat_assistant, + validate_assistant_checkpoint, +) +from executorch.examples.models.gemma4.webgpu_artifact_manifest import ( + create_mtp_manifest, + validate_export_identity, + validate_mtp_manifest, + WEBGPU_EXTERNAL_CONSTANTS_MAX_DATA_BYTES, +) +from executorch.examples.models.gemma4.webgpu_partitioner import ( + build_webgpu_partitioner, + build_webgpu_transform_passes, + rewrite_certified_unique_scatter, +) + + +def _load_target( + checkpoint: Path, + *, + max_seq_len: int, + text_quantize: str, + group_size: int, +) -> Any: + from executorch.examples.models.gemma4.quant_utils import ( + apply_embedding_quantization, + apply_linear_quantization, + parse_quantize, + ) + from executorch.examples.models.gemma4.text_decoder.gemma4_config import ( + Gemma4Config, + ) + from executorch.examples.models.gemma4.text_decoder.gemma4_model import Gemma4Model + + config = Gemma4Config.from_config("e2b") + config.use_kv_cache = True + config.max_seq_len = max_seq_len + config.enable_dynamic_shape = True + config.use_custom_sdpa = True + model = Gemma4Model( + config=config, + checkpoint_path=str(checkpoint.resolve()), + dtype=torch.float32, + ).get_eager_model() + linear_quant, embedding_quant = parse_quantize(text_quantize) + if embedding_quant: + model = apply_embedding_quantization(model, embedding_quant).eval() + if linear_quant: + model = apply_linear_quantization( + model, linear_quant, group_size=group_size + ).eval() + return model.eval() + + +def build_k2_round_program( + target_checkpoint: Path, + assistant_checkpoint: Path, + *, + max_seq_len: int = 8960, + max_input_len: int = 512, + text_quantize: str = "8da4w+emb4", + assistant_quantize: str = "8da4w", + assistant_lm_head_bits: int = 4, + group_size: int = 128, +) -> torch.export.ExportedProgram: + if max_seq_len < 514: + raise ValueError("Gemma 4 MTP export requires max_seq_len >= 514") + if max_input_len < 3 or max_input_len > max_seq_len: + raise ValueError("invalid Gemma 4 MTP max_input_len") + if text_quantize != "8da4w+emb4": + raise ValueError("Gemma 4 MTP target requires 8da4w+emb4 quantization") + if assistant_quantize != "8da4w" or assistant_lm_head_bits != 4: + raise ValueError("Gemma 4 MTP assistant requires 8da4w with a 4-bit LM head") + if group_size != 128: + raise ValueError("Gemma 4 MTP export requires quantization group size 128") + target_checkpoint_evidence = validate_export_identity(target_checkpoint) + validate_assistant_checkpoint(assistant_checkpoint) + target: Any = _load_target( + target_checkpoint / "model.safetensors", + max_seq_len=max_seq_len, + text_quantize=text_quantize, + group_size=group_size, + ) + assistant: Any = load_qat_assistant( + assistant_checkpoint, + max_donor_len=max_seq_len, + lm_head_bits=assistant_lm_head_bits, + quantize_backbone=assistant_quantize, + ) + text_model: Any = target.model + program = compose_k2_round_program( + text_model=text_model, + embed_tokens=text_model.self_decoder.embed_tokens, + assistant=assistant, + hidden_size=text_model.config.hidden_size, + max_seq_len=max_seq_len, + embed_scale=text_model.self_decoder.embed_scale, + max_input_len=max_input_len, + max_donor_len=max_seq_len, + ) + qat_selection_evidence = assistant._webgpu_qat_selection_evidence + rewrites = rewrite_certified_unique_scatter( + program, + assistant.assistant.masked_embedding.token_ordering, + expected_chains=2, + ) + if rewrites != 2: + raise ValueError(f"Gemma 4 MTP expected two scatter rewrites, found {rewrites}") + k2_abi_evidence = program.graph_module.meta.get("gemma4K2Abi") + if not isinstance(k2_abi_evidence, dict): + raise ValueError("K=2 composition lacks exact ABI evidence") + + from executorch.exir.program._program import _transform + from executorch.extension.llm.export.export_passes import ( + ReplaceSDPAWithCustomSDPAPass, + ) + + transformed = _transform(program, ReplaceSDPAWithCustomSDPAPass()) + transformed.graph_module.meta["gemma4QATSelectionEvidence"] = qat_selection_evidence + transformed.graph_module.meta["gemma4K2Abi"] = k2_abi_evidence + transformed.graph_module.meta["gemma4TargetCheckpointEvidence"] = dict( + target_checkpoint_evidence + ) + return transformed + + +def _lower_k2_round( + program: torch.export.ExportedProgram, + *, + external_constants_max_data_bytes: int, + text_quantize: str, +) -> tuple[Any, dict[str, object]]: + from executorch.exir import EdgeCompileConfig, to_edge_transform_and_lower + from executorch.exir.capture._config import ExecutorchBackendConfig + from executorch.exir.passes import MemoryPlanningPass + from executorch.exir.passes.sym_shape_eval_pass import ( + ConstraintBasedSymShapeEvalPass, + ) + + compile_options = { + "alias_buffer_mutations": True, + "external_constants_max_data_bytes": external_constants_max_data_bytes, + "require_dynamic_shapes": True, + } + mtp_transform_passes = build_webgpu_transform_passes(mode="mtp") + edge = to_edge_transform_and_lower( + {"k2_round": program.run_decompositions({})}, + compile_config=EdgeCompileConfig(_check_ir_validity=False), + partitioner={ + "k2_round": [ + build_webgpu_partitioner( + text_quantize=text_quantize, + mode="mtp", + compile_options=compile_options, + ) + ] + }, + transform_passes={"k2_round": mtp_transform_passes}, + ) + edge_census = getattr(mtp_transform_passes[0], "census", None) + if not isinstance(edge_census, dict): + raise ValueError("K=2 lowering did not record the MTP edge census") + lowered = edge.exported_program("k2_round") + delegate_count = 0 + portable_nodes: list[str] = [] + for node in lowered.graph.nodes: + if node.op != "call_function" or node.target is operator.getitem: + continue + target = str(node.target) + if "executorch_call_delegate" in target: + delegate_count += 1 + else: + portable_nodes.append(target) + if delegate_count != 1 or portable_nodes: + raise ValueError( + "K=2 lowering requires one delegate and no portable operators: " + f"delegates={delegate_count}, portable={portable_nodes}" + ) + executorch_program = edge.to_executorch( + ExecutorchBackendConfig( + external_constants=True, + extract_delegate_segments=True, + memory_planning_pass=MemoryPlanningPass(alloc_graph_input=False), + sym_shape_eval_pass=ConstraintBasedSymShapeEvalPass(), + ) + ) + return executorch_program, { + "delegate_count": delegate_count, + "edge": edge_census, + "portable_operator_count": len(portable_nodes), + } + + +def _validate_output_paths(output_path: Path, receipt_path: Path) -> None: + if output_path.suffix != ".pte": + raise ValueError("Gemma 4 MTP output path must end in .pte") + sealed_root = output_path.parent.resolve() + resolved_receipt = receipt_path.resolve() + try: + resolved_receipt.relative_to(sealed_root) + except ValueError: + pass + else: + raise ValueError("Gemma 4 MTP receipt must be outside the sealed artifact root") + + +def export_speculative( # noqa: C901 + target_checkpoint: Path, + assistant_checkpoint: Path, + output_path: Path, + receipt_path: Path, + *, + source_receipt_path: Path | None = None, + max_seq_len: int = 8960, + max_input_len: int = 512, + text_quantize: str = "8da4w+emb4", + assistant_quantize: str = "8da4w", + assistant_lm_head_bits: int = 4, + group_size: int = 128, + external_constants_max_data_bytes: int = (WEBGPU_EXTERNAL_CONSTANTS_MAX_DATA_BYTES), +) -> Path: + _validate_output_paths(output_path, receipt_path) + if (max_seq_len, max_input_len) != (8960, 512): + raise ValueError("production Gemma 4 MTP export requires P512/ctx8960") + if external_constants_max_data_bytes != WEBGPU_EXTERNAL_CONSTANTS_MAX_DATA_BYTES: + raise ValueError( + "production Gemma 4 MTP export requires the reviewed PTD split" + ) + if output_path.exists() or output_path.is_symlink(): + raise ValueError(f"refusing to overwrite existing artifact: {output_path}") + if receipt_path.exists() or receipt_path.is_symlink(): + raise ValueError(f"refusing to overwrite existing artifact: {receipt_path}") + output_path.parent.mkdir(parents=True, exist_ok=True) + receipt_path.parent.mkdir(parents=True, exist_ok=True) + assistant_checkpoint_evidence = validate_assistant_checkpoint(assistant_checkpoint) + program = build_k2_round_program( + target_checkpoint, + assistant_checkpoint, + max_seq_len=max_seq_len, + max_input_len=max_input_len, + text_quantize=text_quantize, + assistant_quantize=assistant_quantize, + assistant_lm_head_bits=assistant_lm_head_bits, + group_size=group_size, + ) + qat_selection_evidence = program.graph_module.meta.get("gemma4QATSelectionEvidence") + if not isinstance(qat_selection_evidence, dict): + raise ValueError("K=2 export lacks QAT selection evidence") + target_checkpoint_evidence = program.graph_module.meta.get( + "gemma4TargetCheckpointEvidence" + ) + if not isinstance(target_checkpoint_evidence, dict): + raise ValueError("K=2 export lacks target checkpoint evidence") + k2_abi_evidence = program.graph_module.meta.get("gemma4K2Abi") + if not isinstance(k2_abi_evidence, dict): + raise ValueError("K=2 export lacks exact ABI evidence") + executorch_program, lowering_evidence = _lower_k2_round( + program, + external_constants_max_data_bytes=external_constants_max_data_bytes, + text_quantize=text_quantize, + ) + tensor_tags = sorted(executorch_program._tensor_data) + if not tensor_tags: + raise ValueError("K=2 export did not produce external tensor data") + for tag in tensor_tags: + if not tag or tag in {".", ".."} or "/" in tag or "\\" in tag: + raise ValueError(f"invalid external tensor-data tag: {tag!r}") + destination = output_path.parent / f"{tag}.ptd" + if destination.exists() or destination.is_symlink(): + raise ValueError(f"refusing to overwrite existing artifact: {destination}") + + with tempfile.TemporaryDirectory( + prefix=f".{output_path.stem}.", dir=output_path.parent.parent + ) as staging_directory, tempfile.TemporaryDirectory( + prefix=f".{receipt_path.stem}.", dir=receipt_path.parent + ) as receipt_staging_directory: + staging = Path(staging_directory) + staged_pte = staging / output_path.name + with staged_pte.open("xb") as output: + executorch_program.write_to_file(output) + executorch_program.write_tensor_data_to_file(str(staging)) + + staged_tensor_paths: list[Path] = [] + for tag in tensor_tags: + path = staging / f"{tag}.ptd" + if path.is_symlink() or not path.is_file() or path.stat().st_size == 0: + raise ValueError(f"missing external tensor data: {path.name}") + staged_tensor_paths.append(path) + if staged_pte.stat().st_size == 0: + raise ValueError("K=2 export produced an empty PTE") + + role_paths: dict[str, Path] = {"pte": staged_pte} + staged_source: Path | None = None + if source_receipt_path is not None: + if source_receipt_path.is_symlink() or not source_receipt_path.is_file(): + raise ValueError("Gemma 4 MTP source receipt must be a regular file") + staged_source = staging / source_receipt_path.name + if staged_source.exists() or staged_source.is_symlink(): + raise ValueError( + "Gemma 4 MTP source receipt basename collides with an artifact" + ) + destination = output_path.parent / staged_source.name + if destination.exists() or destination.is_symlink(): + raise ValueError(f"refusing to overwrite existing artifact: {destination}") + shutil.copyfile(source_receipt_path, staged_source) + role_paths["source"] = staged_source + receipt = create_mtp_manifest(staging, role_paths, staged_tensor_paths) + receipt["evidence"] = { + "assistant_checkpoint": assistant_checkpoint_evidence, + "k2_abi": k2_abi_evidence, + "lowering": lowering_evidence, + "qat_selection": qat_selection_evidence, + "target_checkpoint": target_checkpoint_evidence, + } + validate_mtp_manifest(staging, receipt) + staged_receipt = Path(receipt_staging_directory) / receipt_path.name + staged_receipt.write_text( + json.dumps(receipt, indent=2, sort_keys=True) + "\n", + encoding="utf-8", + ) + + publications = [ + (path, output_path.parent / path.name) for path in staged_tensor_paths + ] + if staged_source is not None: + publications.append( + (staged_source, output_path.parent / staged_source.name) + ) + publications.append((staged_pte, output_path)) + published: list[Path] = [] + try: + for staged, destination in publications: + staged.replace(destination) + published.append(destination) + validate_mtp_manifest(output_path.parent, receipt) + staged_receipt.replace(receipt_path) + published.append(receipt_path) + final_receipt = json.loads(receipt_path.read_text(encoding="utf-8")) + if not isinstance(final_receipt, dict): + raise ValueError("Gemma 4 MTP receipt must be a JSON object") + validate_mtp_manifest(output_path.parent, final_receipt) + except (OSError, ValueError): + for destination in reversed(published): + destination.unlink(missing_ok=True) + raise + return receipt_path + + +def main() -> int: + parser = argparse.ArgumentParser(description="Export Gemma 4 K=2 for WebGPU") + parser.add_argument("--target-checkpoint", type=Path, required=True) + parser.add_argument("--assistant-checkpoint", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--receipt", type=Path, required=True) + parser.add_argument("--source-receipt", type=Path, default=None) + parser.add_argument("--max-seq-len", type=int, default=8960) + parser.add_argument("--max-input-len", type=int, default=512) + args = parser.parse_args() + export_speculative( + args.target_checkpoint, + args.assistant_checkpoint, + args.output, + args.receipt, + source_receipt_path=args.source_receipt, + max_seq_len=args.max_seq_len, + max_input_len=args.max_input_len, + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/examples/models/gemma4/manifests/gemma4_e2b_mtp_webgpu.json b/examples/models/gemma4/manifests/gemma4_e2b_mtp_webgpu.json new file mode 100644 index 00000000000..5710c10cbd2 --- /dev/null +++ b/examples/models/gemma4/manifests/gemma4_e2b_mtp_webgpu.json @@ -0,0 +1,161 @@ +{ + "acquisition": { + "assistant": { + "files": { + "config.json": { + "bytes": 2356, + "sha256": "5d01e9f3f8e969aa8147201a26e849c05446c7c746fa918101ed0622b201db15" + }, + "model.safetensors": { + "bytes": 157565344, + "sha256": "28b11aa1fef73e655107984e0024ed1b149df4b8b36dcb95f27cca603eabc960" + } + }, + "repo_id": "google/gemma-4-E2B-it-qat-q4_0-unquantized-assistant", + "revision": "ebc7e1a211354561464cb82ed6d886792138dcb6" + }, + "target": { + "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": 902378944, + "path": "gemma4_eagle_webgpu_k2.pte", + "role": "pte", + "sha256": "bee968d30b0628d9492ac4a115e14856ca7b12b25932ad3ec49eecec1af01ac4" + }, + { + "bytes": 492608256, + "path": "vulkan_constants_1836442febaaa50dec351c5ca511bf3ce2278cfcd985ae51a7edb9bfbd1c2ad5.ptd", + "role": "ptd", + "sha256": "70a884869849e0620f8558c8ea8ec588464e06a2a79e626f5defd70d57782d22" + }, + { + "bytes": 1499599876, + "path": "vulkan_constants_6216f54d076cc6c0f123eb27824ed18bbd2ab8b730b3e0405582d4273a16f4fb.ptd", + "role": "ptd", + "sha256": "f622f44e0fa06aa8b53b58da4ee8f56b55a4f36d9f68c91683780e10582a7e0c" + }, + { + "bytes": 1442537472, + "path": "vulkan_constants_965014fd4d8400017a16e2582285050b23961e7a9cf0f1b9a9073c957a24e992.ptd", + "role": "ptd", + "sha256": "f6c14986727ea812c6863823627e93204aff852d9ece6b49cc1d843b6504c47c" + } + ], + "export": { + "assistant_calls_per_round": 2, + "assistant_lm_head_bits": 4, + "assistant_quantization": "8da4w", + "backend": "webgpu", + "donor_length": { + "max": 8960, + "min": 2 + }, + "max_input_len": 512, + "max_seq_len": 8960, + "methods": [ + "k2_round" + ], + "selection": { + "centroid_top_k": 32, + "logical_token_ordering_shape": [ + 2048, + 128 + ], + "num_centroids": 2048, + "raw_token_ordering_shape": [ + 262144 + ], + "selected_token_count": 4096, + "tokens_per_centroid": 128 + }, + "speculation_k": 2, + "target_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 + }, + "assistant": { + "architecture": "Gemma4AssistantForCausalLM", + "backboneHiddenSize": 1536, + "hiddenSize": 256, + "modelType": "gemma4_assistant", + "numHiddenLayers": 4, + "vocabSize": 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_1836442febaaa50dec351c5ca511bf3ce2278cfcd985ae51a7edb9bfbd1c2ad5.ptd", + "vulkan_constants_6216f54d076cc6c0f123eb27824ed18bbd2ab8b730b3e0405582d4273a16f4fb.ptd", + "vulkan_constants_965014fd4d8400017a16e2582285050b23961e7a9cf0f1b9a9073c957a24e992.ptd" + ], + "schema_version": 1 +} diff --git a/examples/models/gemma4/mtp_qat_contract.py b/examples/models/gemma4/mtp_qat_contract.py new file mode 100644 index 00000000000..7aee265e1de --- /dev/null +++ b/examples/models/gemma4/mtp_qat_contract.py @@ -0,0 +1,48 @@ +# 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 + +import hashlib +from typing import Any + +import torch + + +OFFICIAL_QAT_NUM_CENTROIDS: int = 2048 +OFFICIAL_QAT_TOKENS_PER_CENTROID: int = 128 +OFFICIAL_QAT_CENTROID_TOP_K: int = 32 +OFFICIAL_QAT_SELECTED_TOKEN_COUNT: int = ( + OFFICIAL_QAT_CENTROID_TOP_K * OFFICIAL_QAT_TOKENS_PER_CENTROID +) + + +def validate_qat_token_ordering(ordering: torch.Tensor) -> dict[str, Any]: + if ordering.numel() != ( + OFFICIAL_QAT_NUM_CENTROIDS * OFFICIAL_QAT_TOKENS_PER_CENTROID + ): + raise ValueError("QAT token ordering must contain 262144 entries") + if ordering.dtype not in (torch.int32, torch.int64): + raise ValueError("QAT token ordering must use an integer dtype") + + raw_shape = list(ordering.shape) + canonical = ordering.detach().to(dtype=torch.int64, device="cpu").reshape(-1) + sorted_values = torch.sort(canonical).values + expected = torch.arange(canonical.numel(), dtype=torch.int64) + if not torch.equal(sorted_values, expected): + raise ValueError("QAT token ordering must be an exact permutation") + + digest = hashlib.sha256(canonical.numpy().tobytes()).hexdigest() + return { + "max": int(canonical.max().item()), + "min": int(canonical.min().item()), + "numel": canonical.numel(), + "permutationExact": True, + "rawShape": raw_shape, + "sha256": digest, + "shape": [OFFICIAL_QAT_NUM_CENTROIDS, OFFICIAL_QAT_TOKENS_PER_CENTROID], + "uniqueCount": canonical.numel(), + } diff --git a/examples/models/gemma4/targets.bzl b/examples/models/gemma4/targets.bzl index 450c8fb676d..23297e95e87 100644 --- a/examples/models/gemma4/targets.bzl +++ b/examples/models/gemma4/targets.bzl @@ -20,6 +20,7 @@ def define_webgpu_python_targets(): fbcode_target(_kind = runtime.python_library, name = "webgpu_support", srcs = [ + "mtp_qat_contract.py", "webgpu_artifact_manifest.py", "webgpu_partitioner.py", ], diff --git a/examples/models/gemma4/webgpu_artifact_manifest.py b/examples/models/gemma4/webgpu_artifact_manifest.py index 8644d80de21..fbd7ff5587b 100644 --- a/examples/models/gemma4/webgpu_artifact_manifest.py +++ b/examples/models/gemma4/webgpu_artifact_manifest.py @@ -117,6 +117,55 @@ "[ExecuTorch][WebGPU] Add Gemma 4 plain export and artifact contract", "[ExecuTorch][WebGPU] Add plain Gemma 4 source-closure tests", ) +MTP_EXPORT_CONTRACT: dict[str, object] = { + "assistant_calls_per_round": 2, + "assistant_lm_head_bits": 4, + "assistant_quantization": "8da4w", + "backend": "webgpu", + "donor_length": {"min": 2, "max": 8960}, + "max_input_len": 512, + "max_seq_len": 8960, + "methods": ["k2_round"], + "selection": { + "centroid_top_k": 32, + "num_centroids": 2048, + "raw_token_ordering_shape": [262144], + "logical_token_ordering_shape": [2048, 128], + "selected_token_count": 4096, + "tokens_per_centroid": 128, + }, + "speculation_k": 2, + "target_quantization": "8da4w+emb4", +} +MTP_SOURCE_CONFIG: dict[str, object] = { + "path": "config/e2b_config.json", + "sha256": SOURCE_CONFIG_SHA256, +} +MTP_ACCEPTED_PROVENANCE: dict[str, object] = { + "artifact_status": "accepted_behavior_oracle", + "source_closure": "pending_final_source_rebuild", +} +MTP_PENDING_SOURCE_PROVENANCE: dict[str, object] = { + "artifact_status": "generated_from_current_source", + "source_closure": "pending_final_source_receipt", +} +MTP_SOURCE_VERIFIED_PROVENANCE: dict[str, object] = { + "artifact_status": "generated_from_current_source", + "source_closure": "source_verified", +} +MTP_PENDING_SOURCE_CLOSURES: frozenset[str] = frozenset( + { + str(MTP_ACCEPTED_PROVENANCE["source_closure"]), + str(MTP_PENDING_SOURCE_PROVENANCE["source_closure"]), + } +) +MTP_EDGE_CENSUS: dict[str, int] = { + "custom_scatter": 2, + "gemma_sdpa": 43, + "generic_scatter": 0, + "legacy_custom_sdpa": 0, + "topk": 2, +} def _source_config_path() -> Path: @@ -873,6 +922,9 @@ def main(argv: Sequence[str] | None = None) -> int: validate = subparsers.add_parser("validate") validate.add_argument("--root", type=Path, required=True) validate.add_argument("--manifest", type=Path, required=True) + validate_mtp = subparsers.add_parser("validate-mtp") + validate_mtp.add_argument("--root", type=Path, required=True) + validate_mtp.add_argument("--manifest", type=Path, required=True) create_source = subparsers.add_parser("create-source-manifest") create_source.add_argument("--fbsource-root", type=Path, required=True) create_source.add_argument("--oss-root", type=Path, required=True) diff --git a/examples/models/gemma4/webgpu_partitioner.py b/examples/models/gemma4/webgpu_partitioner.py index 8e99733fdfe..3339f38fb77 100644 --- a/examples/models/gemma4/webgpu_partitioner.py +++ b/examples/models/gemma4/webgpu_partitioner.py @@ -8,10 +8,14 @@ """Gemma 4-specific, instance-scoped WebGPU partitioning.""" +import operator +from collections.abc import Mapping from functools import lru_cache -from typing import Callable, List, Optional, Tuple +from typing import Any, Callable, List, Literal, Optional, Tuple +import executorch.backends.vulkan.custom_ops_lib # noqa: F401 import executorch.backends.vulkan.patterns as vk_patterns +import executorch.backends.vulkan.utils as vk_utils import torch from executorch.backends.vulkan.op_registry import get_op_features, OpFeatures, OpKey @@ -20,6 +24,13 @@ create_hf_rotary_emb_single_custom_op, HfRotaryEmbeddingSinglePattern, ) +from executorch.examples.models.gemma4.mtp_qat_contract import ( + OFFICIAL_QAT_CENTROID_TOP_K, + OFFICIAL_QAT_NUM_CENTROIDS, + OFFICIAL_QAT_SELECTED_TOKEN_COUNT, + OFFICIAL_QAT_TOKENS_PER_CENTROID, + validate_qat_token_ordering, +) 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 @@ -29,6 +40,24 @@ _EXPECTED_GEMMA4_SDPA_COUNT = 35 _EXPECTED_SINGLE_HF_ROPE_COUNT = 20 +_EXPECTED_ASSISTANT_SDPA_COUNT = 8 +_EXPECTED_ASSISTANT_SINGLE_HF_ROPE_COUNT = 8 +_OFFICIAL_TOPK_INPUT_SHAPE: tuple[int, ...] = (1, 1, OFFICIAL_QAT_NUM_CENTROIDS) +_OFFICIAL_TOPK_OUTPUT_SHAPE: tuple[int, ...] = ( + 1, + 1, + OFFICIAL_QAT_CENTROID_TOP_K, +) +_OFFICIAL_SCATTER_OUTPUT_SHAPE: tuple[int, ...] = ( + 1, + 1, + OFFICIAL_QAT_NUM_CENTROIDS * OFFICIAL_QAT_TOKENS_PER_CENTROID, +) +_OFFICIAL_SCATTER_UPDATE_SHAPE: tuple[int, ...] = ( + 1, + 1, + OFFICIAL_QAT_SELECTED_TOKEN_COUNT, +) def _single_hf_rope_features() -> OpFeatures: @@ -160,24 +189,574 @@ def call(self, exported_program: ExportedProgram) -> ExportedProgramPassResult: return ExportedProgramPassResult(exported_program, True) +def _argument(node: torch.fx.Node, index: int, name: str, default: Any) -> Any: + if len(node.args) > index: + return node.args[index] + return node.kwargs.get(name, default) + + +def _tensor_meta(value: object) -> torch.Tensor | None: + if not isinstance(value, torch.fx.Node): + return None + tensor = value.meta.get("val") + return tensor if isinstance(tensor, torch.Tensor) else None + + +def _tensor_shape(value: object) -> tuple[int, ...] | None: + tensor = _tensor_meta(value) + return tuple(tensor.shape) if tensor is not None else None + + +def _require_tensor( + value: object, + label: str, + shape: tuple[int, ...], + dtype: torch.dtype, +) -> torch.fx.Node: + tensor = _tensor_meta(value) + if not isinstance(value, torch.fx.Node) or tensor is None: + raise ValueError(f"MTP residual {label} is not a tensor node") + if tuple(tensor.shape) != shape or tensor.dtype != dtype: + raise ValueError( + f"MTP residual {label} mismatch: " + f"shape={tuple(tensor.shape)}, dtype={tensor.dtype}" + ) + return value + + +def _semantic_call_users(node: torch.fx.Node) -> list[torch.fx.Node]: + return [ + user + for user in node.users + if user.op == "call_function" + and user.target != torch.ops.aten._assert_tensor_metadata.default + ] + + +def _only_call_user(node: torch.fx.Node, target: object, label: str) -> torch.fx.Node: + users = _semantic_call_users(node) + if len(users) != 1 or users[0].target != target: + raise ValueError(f"MTP residual {label} provenance mismatch") + return users[0] + + +def _is_integer_tensor(value: torch.Tensor) -> bool: + return value.dtype in (torch.int32, torch.int64) + + +def _is_official_qat_topk(node: torch.fx.Node) -> bool: + input_value = _tensor_meta(node.args[0]) + outputs = node.meta.get("val") + if ( + input_value is None + or input_value.dtype != torch.float32 + or tuple(input_value.shape) != _OFFICIAL_TOPK_INPUT_SHAPE + or not isinstance(outputs, (tuple, list)) + or len(outputs) != 2 + or not all(isinstance(value, torch.Tensor) for value in outputs) + ): + return False + values, indices = outputs + return bool( + values.dtype == torch.float32 + and tuple(values.shape) == _OFFICIAL_TOPK_OUTPUT_SHAPE + and _is_integer_tensor(indices) + and tuple(indices.shape) == _OFFICIAL_TOPK_OUTPUT_SHAPE + and _argument(node, 1, "k", None) == OFFICIAL_QAT_CENTROID_TOP_K + and _argument(node, 2, "dim", -1) == -1 + and _argument(node, 3, "largest", True) is True + and _argument(node, 4, "sorted", True) is True + ) + + +def _is_official_qat_unique_scatter(node: torch.fx.Node) -> bool: + if len(node.args) < 4 or _argument(node, 1, "dim", None) != -1: + return False + output = _tensor_meta(node.args[0]) + index = _tensor_meta(node.args[2]) + source = _tensor_meta(node.args[3]) + result = node.meta.get("val") + return bool( + output is not None + and index is not None + and source is not None + and isinstance(result, torch.Tensor) + and output.dtype == torch.float32 + and tuple(output.shape) == _OFFICIAL_SCATTER_OUTPUT_SHAPE + and _is_integer_tensor(index) + and tuple(index.shape) == _OFFICIAL_SCATTER_UPDATE_SHAPE + and source.dtype == torch.float32 + and tuple(source.shape) == _OFFICIAL_SCATTER_UPDATE_SHAPE + and result.dtype == torch.float32 + and tuple(result.shape) == _OFFICIAL_SCATTER_OUTPUT_SHAPE + ) + + +def _lifted_tensor_value( + program: torch.export.ExportedProgram, node: torch.fx.Node +) -> torch.Tensor | None: + if node.op != "placeholder": + return None + signature = program.graph_signature + target = ( + signature.inputs_to_parameters.get(node.name) + or signature.inputs_to_buffers.get(node.name) + or signature.inputs_to_lifted_tensor_constants.get(node.name) + ) + if target is None: + return None + value = program.state_dict.get(target) + if value is None: + value = program.constants.get(target) + return value if isinstance(value, torch.Tensor) else None + + +def _validate_ordering_constant( + program: torch.export.ExportedProgram, + node: torch.fx.Node, + expected_ordering: torch.Tensor, +) -> None: + ordering = _lifted_tensor_value(program, node) + expected = ( + expected_ordering.detach() + .to(torch.int64) + .reshape( + OFFICIAL_QAT_NUM_CENTROIDS, + OFFICIAL_QAT_TOKENS_PER_CENTROID, + ) + ) + if ( + ordering is None + or ordering.dtype != torch.float32 + or tuple(ordering.shape) != tuple(expected.shape) + or not torch.equal(ordering.detach().to(torch.int64).cpu(), expected.cpu()) + ): + raise ValueError("MTP residual token-ordering identity mismatch") + + +def _validate_output_template( + program: torch.export.ExportedProgram, node: torch.fx.Node +) -> None: + template = _lifted_tensor_value(program, node) + if ( + template is None + or template.dtype != torch.float32 + or tuple(template.shape) != _OFFICIAL_SCATTER_OUTPUT_SHAPE + or not bool( + torch.all(template.detach().cpu() == torch.finfo(torch.float32).min) + ) + ): + raise ValueError("MTP residual output-template identity mismatch") + + +def _validate_selected_embedding_chain( + convert: torch.fx.Node, +) -> tuple[torch.fx.Node, torch.fx.Node]: + views = [ + user + for user in _semantic_call_users(convert) + if user.target == torch.ops.aten.view.default + ] + index_view = next( + (user for user in views if _tensor_shape(user) == (1, 1, 4096)), + None, + ) + flat_view = next( + (user for user in views if _tensor_shape(user) == (4096,)), + None, + ) + if index_view is None or flat_view is None or len(views) != 2: + raise ValueError("MTP residual selected-token view topology mismatch") + selected = _only_call_user( + flat_view, + torch.ops.quantized_decomposed.embedding_4bit.dtype, + "selected-token embedding", + ) + if ( + len(selected.args) != 6 + or selected.args[2] is not None + or selected.args[3:5] != (-8, 7) + or selected.args[5] is not flat_view + or selected.kwargs.get("dtype") != torch.float32 + ): + raise ValueError("MTP residual selected-token embedding ABI mismatch") + _require_tensor( + selected, + "selected-token embedding", + (OFFICIAL_QAT_SELECTED_TOKEN_COUNT, 256), + torch.float32, + ) + return index_view, selected + + +def _validate_scatter_source_chain( + scatter: torch.fx.Node, + selected: torch.fx.Node, + topk_input: torch.fx.Node, +) -> None: + source = _require_tensor( + scatter.args[3], + "scatter source", + _OFFICIAL_SCATTER_UPDATE_SHAPE, + torch.float32, + ) + if source.target != torch.ops.aten.squeeze.dim or _semantic_call_users(source) != [ + scatter + ]: + raise ValueError("MTP residual scatter-source provenance mismatch") + matmul = source.args[0] if source.args else None + if ( + not isinstance(matmul, torch.fx.Node) + or matmul.target != torch.ops.aten.matmul.default + ): + raise ValueError("MTP residual scatter-source matmul mismatch") + selected_view = _only_call_user( + selected, torch.ops.aten.view.default, "selected-token view" + ) + selected_transpose = _only_call_user( + selected_view, torch.ops.aten.transpose.int, "selected-token transpose" + ) + if len(matmul.args) != 2 or matmul.args[1] is not selected_transpose: + raise ValueError("MTP residual selected-token matmul linkage mismatch") + query = matmul.args[0] + if ( + not isinstance(query, torch.fx.Node) + or query.target != torch.ops.aten.unsqueeze.default + or not query.args + or not topk_input.args + or query.args[0] is not topk_input.args[0] + ): + raise ValueError("MTP residual score/source hidden-state linkage mismatch") + + +def _certify_scatter_chain( + program: torch.export.ExportedProgram, + scatter: torch.fx.Node, + expected_ordering: torch.Tensor, +) -> tuple[torch.fx.Node, str, str]: + if not _is_official_qat_unique_scatter(scatter): + raise ValueError("MTP export found a non-official scatter contract") + index = _require_tensor( + scatter.args[2], + "scatter index", + _OFFICIAL_SCATTER_UPDATE_SHAPE, + torch.int64, + ) + if index.target != torch.ops.aten.view.default or not index.args: + raise ValueError("MTP residual scatter-index view mismatch") + convert = index.args[0] + if ( + not isinstance(convert, torch.fx.Node) + or convert.target != torch.ops.aten._to_copy.default + or convert.kwargs != {"dtype": torch.int64} + or not convert.args + ): + raise ValueError("MTP residual token-ordering conversion mismatch") + embedding = _require_tensor( + convert.args[0], + "token-ordering embedding", + (1, 1, 32, 128), + torch.float32, + ) + if embedding.target != torch.ops.aten.embedding.default or len(embedding.args) < 2: + raise ValueError("MTP residual token-ordering embedding mismatch") + ordering_node, projection = embedding.args[:2] + if not isinstance(ordering_node, torch.fx.Node) or not isinstance( + projection, torch.fx.Node + ): + raise ValueError("MTP residual ordering/projection is not a graph node") + _validate_ordering_constant(program, ordering_node, expected_ordering) + if projection.target is not operator.getitem or projection.args[1] != 1: + raise ValueError("MTP residual top-k projection mismatch") + topk = projection.args[0] + if not isinstance(topk, torch.fx.Node) or not _is_official_qat_topk(topk): + raise ValueError("MTP residual top-k contract mismatch") + if _semantic_call_users(topk) != [projection]: + raise ValueError("MTP residual top-k consumer mismatch") + topk_input = topk.args[0] + if ( + not isinstance(topk_input, torch.fx.Node) + or topk_input.target != torch.ops.aten.linear.default + or _semantic_call_users(topk_input) != [topk] + ): + raise ValueError("MTP residual top-k producer provenance mismatch") + index_view, selected = _validate_selected_embedding_chain(convert) + if index_view is not index: + raise ValueError("MTP residual scatter-index identity mismatch") + _validate_scatter_source_chain(scatter, selected, topk_input) + template = scatter.args[0] + if not isinstance(template, torch.fx.Node): + raise ValueError("MTP residual output template is not a graph node") + _validate_output_template(program, template) + return topk, ordering_node.name, template.name + + +def mtp_extra_op_features() -> Mapping[OpKey, OpFeatures]: + features = _extra_op_features() + features.update( + { + exir_ops.edge.aten.topk.default: OpFeatures( + inputs_dtypes=[vk_utils.FP_T], + outputs_dtypes=[vk_utils.FP_T, vk_utils.INT_T], + inputs_storage=[vk_utils.CONTIGUOUS_BUFFER], + outputs_storage=[ + vk_utils.CONTIGUOUS_BUFFER, + vk_utils.CONTIGUOUS_BUFFER, + ], + supports_resize=True, + are_node_inputs_supported_fn=_is_official_qat_topk, + ), + exir_ops.edge.et_vk.scatter_src_unique.default: OpFeatures( + inputs_dtypes=[ + vk_utils.FP_T, + vk_utils.NONE_T, + vk_utils.INT_T, + vk_utils.FP_T, + ], + outputs_dtypes=[vk_utils.FP_T], + inputs_storage=[ + vk_utils.CONTIGUOUS_BUFFER, + vk_utils.NO_STORAGE, + vk_utils.CONTIGUOUS_BUFFER, + vk_utils.CONTIGUOUS_BUFFER, + ], + outputs_storage=[vk_utils.CONTIGUOUS_BUFFER], + supports_resize=True, + are_node_inputs_supported_fn=_is_official_qat_unique_scatter, + ), + } + ) + return features + + +def rewrite_certified_unique_scatter( + program: torch.export.ExportedProgram, + token_ordering: torch.Tensor, + *, + expected_chains: int = 2, +) -> int: + validate_qat_token_ordering(token_ordering) + topk_nodes = [ + node + for node in program.graph.nodes + if node.op == "call_function" and node.target == torch.ops.aten.topk.default + ] + scatter_nodes = [ + node + for node in program.graph.nodes + if node.op == "call_function" and node.target == torch.ops.aten.scatter.src + ] + if ( + expected_chains <= 0 + or len(topk_nodes) != expected_chains + or len(scatter_nodes) != expected_chains + ): + raise ValueError( + "MTP export residual topology mismatch: " + f"topk={len(topk_nodes)}, scatter={len(scatter_nodes)}, " + f"expected={expected_chains}" + ) + certified = [ + _certify_scatter_chain(program, node, token_ordering) for node in scatter_nodes + ] + if ( + {record[0] for record in certified} != set(topk_nodes) + or len({record[1] for record in certified}) != 1 + or len({record[2] for record in certified}) != 1 + ): + raise ValueError("MTP export residual-chain identity mismatch") + for node in scatter_nodes: + node.target = torch.ops.et_vk.scatter_src_unique.default + node.meta["gemma4_mtp_unique_scatter_certified"] = True + program.graph_module.recompile() + return len(scatter_nodes) + + +def _node_has_module_fragment(node: torch.fx.Node, fragment: str) -> bool: + stack = node.meta.get("nn_module_stack") or {} + for entry in stack.values(): + path = entry[0] if isinstance(entry, tuple) and entry else str(entry) + if fragment in str(path): + return True + return False + + +def _replace_mtp_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, + ) + nodes = [ + node + for node in graph_module.graph.nodes + if node.target == exir_ops.edge.et_vk.apply_rotary_emb_hf_single.default + ] + target_count = sum(_node_has_module_fragment(node, "target") for node in nodes) + assistant_count = sum( + _node_has_module_fragment(node, "assistant") for node in nodes + ) + if ( + target_count != _EXPECTED_SINGLE_HF_ROPE_COUNT + or assistant_count != _EXPECTED_ASSISTANT_SINGLE_HF_ROPE_COUNT + or len(nodes) != target_count + assistant_count + ): + raise ValueError( + "Gemma4 MTP single-HF-RoPE scope mismatch: " + f"target={target_count}, assistant={assistant_count}, total={len(nodes)}" + ) + + +def _rewrite_mtp_sdpa(exported_program: ExportedProgram) -> None: + graph_module = exported_program.graph_module + target_nodes: list[torch.fx.Node] = [] + assistant_nodes: list[torch.fx.Node] = [] + unscoped_nodes: list[torch.fx.Node] = [] + for node in graph_module.graph.nodes: + if node.target != exir_ops.edge.llama.custom_sdpa.default: + continue + if _node_has_module_fragment(node, "target"): + target_nodes.append(node) + elif _node_has_module_fragment(node, "assistant"): + assistant_nodes.append(node) + else: + unscoped_nodes.append(node) + if len(node.args) != 8 or node.kwargs: + raise ValueError("Gemma4 MTP custom SDPA must use the 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 MTP custom SDPA is not WebGPU-compatible") + if ( + len(target_nodes) != _EXPECTED_GEMMA4_SDPA_COUNT + or len(assistant_nodes) != _EXPECTED_ASSISTANT_SDPA_COUNT + or unscoped_nodes + ): + raise ValueError( + "Gemma4 MTP SDPA scope mismatch: " + f"target={len(target_nodes)}, assistant={len(assistant_nodes)}, " + f"unscoped={len(unscoped_nodes)}" + ) + for node in [*target_nodes, *assistant_nodes]: + node.target = exir_ops.edge.et_vk.gemma4_sdpa.default + graph_module.recompile() + + +def _validate_mtp_edge_census( + exported_program: ExportedProgram, +) -> dict[str, int]: + counts = { + "custom_scatter": 0, + "gemma_sdpa": 0, + "generic_scatter": 0, + "legacy_custom_sdpa": 0, + "topk": 0, + } + targets = { + exir_ops.edge.et_vk.scatter_src_unique.default: "custom_scatter", + exir_ops.edge.et_vk.gemma4_sdpa.default: "gemma_sdpa", + exir_ops.edge.aten.scatter.src: "generic_scatter", + exir_ops.edge.llama.custom_sdpa.default: "legacy_custom_sdpa", + exir_ops.edge.aten.topk.default: "topk", + } + for node in exported_program.graph.nodes: + label = targets.get(node.target) + if label is not None: + counts[label] += 1 + expected = { + "custom_scatter": 2, + "gemma_sdpa": 43, + "generic_scatter": 0, + "legacy_custom_sdpa": 0, + "topk": 2, + } + if counts != expected: + raise ValueError(f"Gemma4 MTP edge census mismatch: {counts}") + return counts + + +class Gemma4MTPWebGPURewritePass(ExportedProgramPassBase): + """Fixed MTP edge rewrites that run before metadata-only partitioning.""" + + def __init__(self) -> None: + super().__init__() + self.census: Optional[dict[str, int]] = None + + def call(self, exported_program: ExportedProgram) -> ExportedProgramPassResult: + _replace_mtp_single_hf_rope(exported_program) + _rewrite_mtp_sdpa(exported_program) + self.census = _validate_mtp_edge_census(exported_program) + exported_program.graph_module.meta["gemma4MTPEdgeCensus"] = dict(self.census) + return ExportedProgramPassResult(exported_program, True) + + +def _mtp_webgpu_allowlist() -> list[OpKey]: + additions = [ + exir_ops.edge.aten.bmm.default, + exir_ops.edge.aten.embedding.default, + exir_ops.edge.aten.mm.default, + exir_ops.edge.aten.eq.Scalar, + exir_ops.edge.aten.sub.Tensor, + exir_ops.edge.aten.where.self, + exir_ops.edge.aten.topk.default, + exir_ops.edge.et_vk.scatter_src_unique.default, + ] + return list(dict.fromkeys([*_webgpu_allowlist(), *additions])) + + class Gemma4WebGPUPartitioner(Partitioner): """Vulkan serialization restricted to Gemma 4 WebGPU capabilities.""" - def __init__(self, text_quantize: str) -> None: + def __init__( + self, + text_quantize: str, + *, + mode: Literal["plain", "mtp"] = "plain", + compile_options: dict[str, Any] | None = None, + ) -> None: + if mode not in ("plain", "mtp"): + raise ValueError(f"invalid Gemma4 WebGPU partitioner mode: {mode!r}") if "emb8" in text_quantize: raise ValueError( "WebGPU cannot delegate emb8; use emb4 (for example, 8da4w+emb4)" ) + if mode == "mtp" and text_quantize != "8da4w+emb4": + raise ValueError("Gemma4 MTP WebGPU requires 8da4w+emb4") + if mode == "plain" and compile_options is not None: + raise ValueError("plain Gemma4 WebGPU does not accept option overrides") + options = { + "external_constants_max_data_bytes": ( + WEBGPU_EXTERNAL_CONSTANTS_MAX_DATA_BYTES + ), + "require_dynamic_shapes": True, + "skip_bool_tensors": mode == "plain", + } + if mode == "mtp": + options["alias_buffer_mutations"] = True + if compile_options is not None: + for key, value in compile_options.items(): + if key in options and value != options[key]: + raise ValueError( + f"Gemma4 MTP WebGPU cannot override {key}: {value!r}" + ) + options[key] = value 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(), + options, + operator_allowlist=( + _webgpu_allowlist() if mode == "plain" else _mtp_webgpu_allowlist() + ), + extra_op_features=( + _extra_op_features() if mode == "plain" else mtp_extra_op_features() + ), ) def ops_to_not_decompose(self, ep: ExportedProgram) -> Tuple[ @@ -193,10 +772,25 @@ def partition(self, exported_program: ExportedProgram) -> PartitionResult: return self._inner.partition(exported_program) -def build_webgpu_partitioner(text_quantize: str) -> Gemma4WebGPUPartitioner: - return Gemma4WebGPUPartitioner(text_quantize) +def build_webgpu_partitioner( + text_quantize: str, + *, + mode: Literal["plain", "mtp"] = "plain", + compile_options: dict[str, Any] | None = None, +) -> Gemma4WebGPUPartitioner: + return Gemma4WebGPUPartitioner( + text_quantize, + mode=mode, + compile_options=compile_options, + ) -def build_webgpu_transform_passes() -> List[ExportedProgramPassBase]: - """Edge transform passes the WebGPU text decoder must run pre-partition.""" - return [_Gemma4WebGPURewritePass()] +def build_webgpu_transform_passes( + mode: Literal["plain", "mtp"] = "plain", +) -> List[ExportedProgramPassBase]: + """Return fixed edge transforms for the requested Gemma4 WebGPU mode.""" + if mode == "plain": + return [_Gemma4WebGPURewritePass()] + if mode == "mtp": + return [Gemma4MTPWebGPURewritePass()] + raise ValueError(f"invalid Gemma4 WebGPU transform mode: {mode!r}")