diff --git a/backends/qualcomm/CMakeLists.txt b/backends/qualcomm/CMakeLists.txt index a74d6546bb7..70f38a180fd 100644 --- a/backends/qualcomm/CMakeLists.txt +++ b/backends/qualcomm/CMakeLists.txt @@ -369,6 +369,24 @@ install( RUNTIME DESTINATION ${_qnn_install_dir} ) +# rpcmem I/O buffers for the LLM runner's IOManagers. It registers itself, so +# runners link it whole-archive. +if(EXECUTORCH_BUILD_EXTENSION_LLM_RUNNER) + add_library( + qnn_llm_io_buffer_allocator STATIC + ${QNN_EXECUTORCH_ROOT_DIR}/llm/io_buffer_allocator.cpp + ) + target_link_libraries( + qnn_llm_io_buffer_allocator PUBLIC qnn_executorch_backend executorch_core + extension_module_static extension_tensor + ) + install( + TARGETS qnn_llm_io_buffer_allocator + EXPORT ExecuTorchTargets + ARCHIVE DESTINATION ${_qnn_install_dir} + ) +endif() + if(${CMAKE_SYSTEM_PROCESSOR} MATCHES Hexagon) add_subdirectory( ${QNN_EXECUTORCH_ROOT_DIR}/runtime/backends/direct_mode diff --git a/backends/qualcomm/_passes/convert_linear_to_conv2d.py b/backends/qualcomm/_passes/convert_linear_to_conv2d.py index c2937a6fa60..e092cf93ebc 100644 --- a/backends/qualcomm/_passes/convert_linear_to_conv2d.py +++ b/backends/qualcomm/_passes/convert_linear_to_conv2d.py @@ -145,6 +145,7 @@ def call(self, graph_module: GraphModule): # The set tracks whether the weights associated with a linear node have been preprocessed, # and is used in cases where multiple linear nodes share the same weights. preprocessed_linear_weights_set = set() + calibrated_inputs = {} for node in graph.nodes: if node.target in {self.aten_linear_target, self.edge_linear_target}: @@ -170,28 +171,34 @@ def call(self, graph_module: GraphModule): # rank = 2: Reshape (dim, C) to (1, dim, C, 1), and permute to (1, C, 1, dim) # rank = 3: Reshape (N, dim, C) to (N, dim, C, 1), and permute to (N, C, 1, dim) # rank = 4: Permute (N, H, W, C) to (N, C, H, W) - if rank <= 3: - shape = ( - (1, *input_meta_val.shape, 1) - if rank == 2 - else (*input_meta_val.shape, 1) - ) - cur_meta_val = cur_meta_val.reshape(shape) - reshape_node = self._create_node( + # Linears reading the same input (q/k/v, gate/up) share one layout + # change: each extra copy is an op HTP runs. + if input_node in calibrated_inputs: + permute_node, cur_meta_val = calibrated_inputs[input_node] + else: + if rank <= 3: + shape = ( + (1, *input_meta_val.shape, 1) + if rank == 2 + else (*input_meta_val.shape, 1) + ) + cur_meta_val = cur_meta_val.reshape(shape) + reshape_node = self._create_node( + graph_module, + ops["view"], + (input_node, shape), + cur_meta_val, + ) + + order = (0, 3, 1, 2) if rank == 4 else (0, 2, 3, 1) + cur_meta_val = cur_meta_val.permute(order) + permute_node = self._create_node( graph_module, - ops["view"], - (input_node, shape), + ops["permute"], + (reshape_node, order) if rank <= 3 else (input_node, order), cur_meta_val, ) - - order = (0, 3, 1, 2) if rank == 4 else (0, 2, 3, 1) - cur_meta_val = cur_meta_val.permute(order) - permute_node = self._create_node( - graph_module, - ops["permute"], - (reshape_node, order) if rank <= 3 else (input_node, order), - cur_meta_val, - ) + calibrated_inputs[input_node] = (permute_node, cur_meta_val) # Step 2: reshape weight if weight_placeholder_node.name not in preprocessed_linear_weights_set: diff --git a/backends/qualcomm/builders/node_visitor.py b/backends/qualcomm/builders/node_visitor.py index 02fbe08e8b0..a68b018b9cb 100644 --- a/backends/qualcomm/builders/node_visitor.py +++ b/backends/qualcomm/builders/node_visitor.py @@ -37,6 +37,7 @@ from .utils import ( deduce_dtype, get_parameter, + graph_output_position, is_graph_input, is_graph_output, is_mutable_buffer_input, @@ -465,6 +466,8 @@ def get_tensor_name( # the input order between QNN and the original graph’s forward function may differ. # The `mutbuf_{id}` is utilized for mapping I/O of mutable buffer at runtime. # The `output_` is identified as the graph’s output at runtime to prevent confusion with per_tensor_dump. + # The `output_{id}` is the position among the graph's outputs, sorted on at runtime like `input_{id}`: + # the same passes can also reorder the ops that produce the outputs. # For more info about self.is_qnn_partitioner, check constructor for explanation. if self.is_qnn_partitioner: @@ -482,7 +485,9 @@ def get_tensor_name( ).index(node.name) tensor_name = f"output_mutbuf_{position_index}_{tensor_name}" elif is_graph_output(node, wrapper_idx): - tensor_name = f"output_{tensor_name}" + position = graph_output_position(node, wrapper_idx) + prefix = "output_" if position is None else f"output_{position}_" + tensor_name = f"{prefix}{tensor_name}" # Only add qcom_tensor_name when enable tensor dump. # Only runs in qnn_preprocess (not op validation) since that's when diff --git a/backends/qualcomm/builders/utils.py b/backends/qualcomm/builders/utils.py index c691ef7ffd3..73e56da84c6 100755 --- a/backends/qualcomm/builders/utils.py +++ b/backends/qualcomm/builders/utils.py @@ -141,6 +141,27 @@ def is_graph_output(node: torch.fx.Node, output_index: Optional[int] = None) -> return False +def graph_output_position( + node: torch.fx.Node, output_index: Optional[int] = None +) -> Optional[int]: + """The position of ``node`` (its ``output_index``-th output, for a + multi-output node) among the graph's outputs, or None.""" + output = next(n for n in node.graph.nodes if n.op == "output") + outputs = output.args[0] + outputs = outputs if isinstance(outputs, (list, tuple)) else [outputs] + for position, value in enumerate(outputs): + if value is node: + return position + if ( + isinstance(value, torch.fx.Node) + and getattr(value.target, "__name__", "") == "getitem" + and value.args[0] is node + and (output_index is None or value.args[1] == output_index) + ): + return position + return None + + def is_mutable_buffer_output( tensor: torch.fx.Node, edge_program: torch.export.ExportedProgram ) -> bool: diff --git a/backends/qualcomm/llm/BUCK b/backends/qualcomm/llm/BUCK new file mode 100644 index 00000000000..54a222191c8 --- /dev/null +++ b/backends/qualcomm/llm/BUCK @@ -0,0 +1,16 @@ +load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target", "non_fbcode_target") +oncall("executorch") +# Any targets that should be shared between fbcode and xplat must be defined in +# targets.bzl. This file can contain xplat-only targets. + +load(":targets.bzl", "define_common_targets") + + +non_fbcode_target(_kind = define_common_targets,) + +# !!!! fbcode/executorch/backends/qualcomm/runtime/TARGETS was merged into this file, see https://fburl.com/workplace/xl8l9yuo for more info !!!! + +load(":targets.bzl", "define_common_targets") + + +fbcode_target(_kind = define_common_targets,) diff --git a/backends/qualcomm/llm/io_buffer_allocator.cpp b/backends/qualcomm/llm/io_buffer_allocator.cpp new file mode 100644 index 00000000000..27cc51e3623 --- /dev/null +++ b/backends/qualcomm/llm/io_buffer_allocator.cpp @@ -0,0 +1,43 @@ +/* + * 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. + */ + +// Puts LLM runner I/O buffers (KV cache, mask, logits) in rpcmem, so HTP reads +// and writes them in place. Linking this library is enough: it registers +// itself as the runner's I/O buffer allocator. Programs must be lowered with +// shared buffers enabled for QNN to use the memory without copying. + +#include +#include + +namespace executorch { +namespace backends { +namespace qnn { +namespace { + +using ::executorch::extension::llm::IOBufferAllocator; + +class SharedBufferAllocator final : public IOBufferAllocator { + public: + void* allocate(size_t nbytes, size_t alignment) override { + return QnnExecuTorchAllocCustomMem(nbytes, alignment); + } + void deallocate(void* block) override { + QnnExecuTorchFreeCustomMem(block); + } + void bind(void* block, void* tensor_data) override { + QnnExecuTorchAddCustomMemTensorAddr(tensor_data, block); + } +}; + +const bool registered = ::executorch::extension::llm::set_io_buffer_allocator( + [] { return std::make_unique(); }); + +} // namespace +} // namespace qnn +} // namespace backends +} // namespace executorch diff --git a/backends/qualcomm/llm/targets.bzl b/backends/qualcomm/llm/targets.bzl new file mode 100644 index 00000000000..6e6664345b3 --- /dev/null +++ b/backends/qualcomm/llm/targets.bzl @@ -0,0 +1,26 @@ +load( + "@fbsource//tools/build_defs:default_platform_defs.bzl", + "ANDROID", +) +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") + +def define_common_targets(): + """Defines targets that should be shared between fbcode and xplat. + + The directory containing this targets.bzl file should also contain both + TARGETS and BUCK files that call this function. + """ + + # rpcmem I/O buffers for the LLM runner's IOManagers. Registers itself on + # load, so link it whole. + runtime.cxx_library( + name = "io_buffer_allocator", + srcs = ["io_buffer_allocator.cpp"], + link_whole = True, + platforms = [ANDROID], + deps = [ + "//executorch/backends/qualcomm/runtime:runtime_android_build", + "//executorch/extension/llm/runner/io_manager:io_manager", + ], + visibility = ["PUBLIC"], + ) diff --git a/backends/qualcomm/partition/qnn_partitioner.py b/backends/qualcomm/partition/qnn_partitioner.py index 23016de006b..ec40abf82c4 100644 --- a/backends/qualcomm/partition/qnn_partitioner.py +++ b/backends/qualcomm/partition/qnn_partitioner.py @@ -4,11 +4,15 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. import copy +import inspect import logging -from collections import defaultdict +from collections import defaultdict, OrderedDict from typing import Any, Callable, Dict, List, Optional, Tuple import torch +from executorch.backends.qualcomm._passes.qnn_pass_manager import ( + get_qnn_pass_manager_cls, +) from executorch.backends.qualcomm.builders import node_visitor_manager from executorch.backends.qualcomm.builders.node_visitor import ( PER_CHANNEL_GROUP_ENCODING, @@ -41,6 +45,7 @@ PartitionResult, ) from executorch.exir.backend.utils import tag_constant_data, tag_mutated_buffer +from executorch.exir.pass_manager import PassType from torch.export.exported_program import ExportedProgram from torch.fx.passes.infra.partitioner import Partition from torch.fx.passes.operator_support import OperatorSupportBase @@ -199,6 +204,9 @@ def __init__( skip_node_id_set: set = None, skip_node_op_set: set = None, skip_mutable_buffer: bool = False, + dep_table: Optional[Dict] = None, + passes_job: Optional[OrderedDict] = None, + convert_linear_to_conv2d: bool = False, ): """ Args: @@ -206,7 +214,15 @@ def __init__( skip_node_id_set (set, optional): Set of node IDs to exclude from partitioning. skip_node_op_set (set, optional): Set of OpOverload to exclude from partitioning. skip_mutable_buffer (bool, optional): If True, mutable buffers are not delegated to QNN. + dep_table (Dict, optional): Dependency table for the transformation passes. + passes_job (OrderedDict, optional): Ordered dictionary of transformation passes. + convert_linear_to_conv2d (bool, optional): Convert linear to conv2d where that + improves HTP performance. """ + self.compiler_specs = compiler_specs + self.dep_table = dep_table + self.passes_job = passes_job + self.convert_linear_to_conv2d = convert_linear_to_conv2d self.compiler_specs_snapshot = copy.deepcopy(compiler_specs) self.backend = flatbuffer_to_option( generate_qnn_executorch_option(self.compiler_specs_snapshot) @@ -220,6 +236,82 @@ def __init__( self.skip_node_op_set = set() if skip_node_op_set is None else skip_node_op_set self.skip_mutable_buffer = skip_mutable_buffer + @classmethod + def from_options( + cls, + soc_model: str, + backend: str = "htp", + **options: Any, + ) -> "QnnPartitioner": + """Build a QnnPartitioner from plain options, generating its compiler specs. + + Lets a caller that only knows option names and values (for example an + exporter forwarding a ``dict`` of vendor settings) get a working + partitioner without importing the QNN compiler-spec builders. + + Args: + soc_model (str): Target SoC name, e.g. ``"SM8750"``. + backend (str): QNN backend, one of ``"htp"``, ``"gpu"`` or ``"lpai"``. + **options: Each option is routed by name to the first of these that + accepts it: ``QnnPartitioner.__init__`` (``skip_node_id_set``, + ``passes_job``, ``convert_linear_to_conv2d``, ...), + ``generate_qnn_executorch_compiler_spec`` (``use_mha2sha``, + ``shared_buffer``, ...), then ``generate__compiler_spec`` + (``use_fp16``, ``htp_performance_mode``, ...). An option none of + them accepts raises ``TypeError``. + """ + from executorch.backends.qualcomm.utils import utils as qnn_utils + + backend_spec_fn = getattr(qnn_utils, f"generate_{backend}_compiler_spec", None) + if backend_spec_fn is None: + raise ValueError(f"Unsupported QNN backend {backend!r}.") + + def accepted(fn): + return set(inspect.signature(fn).parameters) - {"self"} + + partitioner_kwargs, compiler_spec_kwargs, backend_kwargs = {}, {}, {} + for name, value in options.items(): + if name in accepted(cls.__init__) - {"compiler_specs"}: + partitioner_kwargs[name] = value + elif name in accepted(qnn_utils.generate_qnn_executorch_compiler_spec) - { + "soc_model", + "backend_options", + }: + compiler_spec_kwargs[name] = value + elif name in accepted(backend_spec_fn): + backend_kwargs[name] = value + else: + raise TypeError( + f"Unknown QnnPartitioner option {name!r} for backend {backend!r}." + ) + + compiler_specs = qnn_utils.generate_qnn_executorch_compiler_spec( + soc_model=qnn_utils.get_soc_to_chipset_map()[soc_model], + backend_options=backend_spec_fn(**backend_kwargs), + **compiler_spec_kwargs, + ) + return cls(compiler_specs, **partitioner_kwargs) + + def transform_for_pre_decomposition( + self, exported_program: ExportedProgram + ) -> ExportedProgram: + """Run the QNN ATen-level passes before default decompositions.""" + return get_qnn_pass_manager_cls(self.backend)().transform_for_export_pipeline( + exported_program + ) + + def transform_passes(self, exported_program: ExportedProgram) -> List[PassType]: + """Run the QNN edge-level passes before partitioning.""" + return get_qnn_pass_manager_cls(self.backend)().get_to_edge_transform_passes( + exported_program, + passes_job=self.passes_job, + dep_table=self.dep_table, + compiler_specs=self.compiler_specs, + skip_node_id_set=self.skip_node_id_set, + skip_node_op_set=self.skip_node_op_set, + convert_linear_to_conv2d=self.convert_linear_to_conv2d, + ) + def generate_partitions( self, edge_program: torch.export.ExportedProgram ) -> List[Any]: diff --git a/backends/qualcomm/quantizer/custom_annotation.py b/backends/qualcomm/quantizer/custom_annotation.py index 54456877c1f..941d1a45f2d 100644 --- a/backends/qualcomm/quantizer/custom_annotation.py +++ b/backends/qualcomm/quantizer/custom_annotation.py @@ -269,9 +269,14 @@ def annotate_matmul_input1(node: Node, is_qat: str): if ( node.op == "call_function" and node.target == torch.ops.aten.matmul.default - and all(arg.op == "call_function" for arg in node.args) + and all( + arg.op == "call_function" + or (arg.op == "placeholder" and arg.meta["val"].is_floating_point()) + for arg in node.args + ) ): - # Only apply custom annotation on Q @ K^T @ V + # Only apply custom annotation on Q @ K^T @ V (activations, or a + # cache that is a graph input), not on matmuls with constants. annotate_matmul(node, quantization_config_16a8w) annotate_matmul_input1(node.args[1], is_qat=is_qat) diff --git a/backends/qualcomm/quantizer/hf_llm_quant_recipe.py b/backends/qualcomm/quantizer/hf_llm_quant_recipe.py new file mode 100644 index 00000000000..20aad8c1fd7 --- /dev/null +++ b/backends/qualcomm/quantizer/hf_llm_quant_recipe.py @@ -0,0 +1,402 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# 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. + +import logging +from abc import ABC +from typing import Optional + +import torch +from executorch.backends.qualcomm.quantizer.custom_annotation import annotate_kv_8bit +from executorch.backends.qualcomm.quantizer.quant_recipe import ( + QuantGranularity, + QuantRecipe, +) +from executorch.backends.qualcomm.quantizer.quantizer import QuantDtype +from torchao.quantization.pt2e import MinMaxObserver + + +class LLMQuantRecipe(ABC): + """ + LLM quantization recipe. + + Subclasses must define ``default_quant_dtype`` and populate ``self.recipe``. + """ + + default_quant_dtype: QuantDtype = None + # Export settings that go with the recipe, as Qualcomm's native flow pairs + # them per model (QnnExportPlan's ``r3``, ``seq_mse_candidates`` and + # ``masked_softmax``). + r3: bool = False + seq_mse_candidates: int = 0 + masked_softmax: bool = True + + def __init__(self): + self.recipe: Optional[QuantRecipe] = None + + if self.default_quant_dtype is None: + raise ValueError("default_quant_dtype must be defined in the recipe.") + + def get_kv_io_bit_width(self) -> int: + if self.default_quant_dtype is None: + return 32 + elif self.default_quant_dtype == QuantDtype.use_8a8w or annotate_kv_8bit in ( + getattr(c, "func", c) for c in self.recipe.custom_quant_annotations + ): + return 8 + else: + # If quantized but not 8a8w or mix_quantization, it has to be 16bit kv io. + return 16 + + def get_logits_output_bit_width(self) -> int: + # We use 16bit logits for all quant config + return 32 if self.default_quant_dtype is None else 16 + + +class Llama3_2_1B_HFQuantRecipe(LLMQuantRecipe): + default_quant_dtype = QuantDtype.use_16a4w + masked_softmax = False + seq_mse_candidates = 50 + # Native's config leaves R3 off; on HF's graph it closes most of the + # perplexity gap to native (24.67 -> 24.13 vs 24.03, same calibration). + r3 = True + + def __init__(self, verbose: bool = False): + super().__init__() + + self.recipe = ( + QuantRecipe( + self.default_quant_dtype, + False, + act_observer=MinMaxObserver, + granularity=QuantGranularity.PER_TENSOR, + verbose=verbose, + note="default with 16bit activation", + ) + .add_node_target( + { + torch.ops.aten.conv2d.default, + }, + QuantDtype.use_16a4w_block, + False, + act_observer=MinMaxObserver, + granularity=QuantGranularity.PER_BLOCK, + extra_kwargs={"block_size": (1, 32, 1, 1)}, + note="Annotate with 16a4w block quantization since these layers are not sensitive.", + ) + .add_regex( + { + r"(?:^|\.)lm_head\.conv", + r"model\.layers\.[0-3]\.mlp\.down_proj\.conv", + }, + QuantDtype.use_16a8w, + False, + act_observer=MinMaxObserver, + granularity=QuantGranularity.PER_CHANNEL, + note="Head and early down_proj are sensitive and should be annotated with 16a8w.", + ) + ) + self.recipe.custom_quant_annotations.append(annotate_kv_8bit) + + +class Qwen2_5_0_5B_HFQuantRecipe(LLMQuantRecipe): + """Ported from Qwen2_5_0_5BQuantRecipe (static_llm_quant_recipe.py:545).""" + + default_quant_dtype = QuantDtype.use_16a4w + r3 = True + + def __init__(self, verbose: bool = False): + super().__init__() + + self.recipe = ( + QuantRecipe( + self.default_quant_dtype, + False, + act_observer=MinMaxObserver, + granularity=QuantGranularity.PER_TENSOR, + verbose=verbose, + note="default with 16bit activation", + ) + .add_node_target( + { + torch.ops.aten.conv2d.default, + }, + QuantDtype.use_16a4w_block, + False, + act_observer=MinMaxObserver, + granularity=QuantGranularity.PER_BLOCK, + extra_kwargs={"block_size": (1, 16, 1, 1)}, + note="Annotate with 16a4w block quantization since these layers are not sensitive.", + ) + .add_regex( + {r"(?:^|\.)lm_head\.conv"}, + QuantDtype.use_16a8w, + False, + act_observer=MinMaxObserver, + granularity=QuantGranularity.PER_CHANNEL, + note="Head is sensitive and should be annotated with 16a8w.", + ) + ) + + +class Qwen2_5_1_5B_HFQuantRecipe(LLMQuantRecipe): + """Ported from Qwen2_5_1_5BQuantRecipe (static_llm_quant_recipe.py:569). + + static_llama's ``output\\.conv`` head is ``model.lm_head.conv`` in HF. + """ + + default_quant_dtype = QuantDtype.use_16a4w + r3 = True + + def __init__(self, verbose: bool = False): + super().__init__() + + self.recipe = ( + QuantRecipe( + self.default_quant_dtype, + False, + act_observer=MinMaxObserver, + granularity=QuantGranularity.PER_TENSOR, + verbose=verbose, + note="default with 16bit activation", + ) + .add_node_target( + { + torch.ops.aten.conv2d.default, + }, + QuantDtype.use_16a4w_block, + False, + act_observer=MinMaxObserver, + granularity=QuantGranularity.PER_BLOCK, + extra_kwargs={"block_size": (1, 16, 1, 1)}, + note="Annotate with 16a4w block quantization since these layers are not sensitive.", + ) + .add_regex( + {r"(?:^|\.)lm_head\.conv"}, + QuantDtype.use_16a8w, + False, + act_observer=MinMaxObserver, + granularity=QuantGranularity.PER_CHANNEL, + note="Head is sensitive and should be annotated with 16a8w.", + ) + ) + + +class Qwen3_0_6B_HFQuantRecipe(LLMQuantRecipe): + """Ported from Qwen3_0_6BQuantRecipe (static_llm_quant_recipe.py:603). + + static_llama's ``layers\\..*\\.feed_forward\\..*w2_conv`` (the down + projection) is ``model.layers.N.mlp.down_proj.conv`` in HF. + """ + + default_quant_dtype = QuantDtype.use_16a4w + seq_mse_candidates = 50 + + def __init__(self, verbose: bool = False): + super().__init__() + + self.recipe = ( + QuantRecipe( + self.default_quant_dtype, + False, + act_observer=MinMaxObserver, + granularity=QuantGranularity.PER_TENSOR, + verbose=verbose, + note="default with 16bit activation", + ) + .add_node_target( + { + torch.ops.aten.conv2d.default, + }, + QuantDtype.use_16a4w_block, + False, + act_observer=MinMaxObserver, + granularity=QuantGranularity.PER_BLOCK, + extra_kwargs={"block_size": (1, 32, 1, 1)}, + note="Annotate with 16a4w block quantization since these layers are not sensitive.", + ) + .add_regex( + { + r"model\.layers\..*\.mlp\.down_proj\.conv", + }, + QuantDtype.use_16a8w, + False, + act_observer=MinMaxObserver, + granularity=QuantGranularity.PER_CHANNEL, + note="down_proj is sensitive and should be annotated with 16a8w.", + ) + ) + + +class Smollm2_HFQuantRecipe(LLMQuantRecipe): + """Ported from Smollm2QuantRecipe (static_llm_quant_recipe.py:676).""" + + default_quant_dtype = QuantDtype.use_16a8w + masked_softmax = False + r3 = True + + def __init__(self, verbose: bool = False): + super().__init__() + + self.recipe = QuantRecipe( + self.default_quant_dtype, + False, + act_observer=MinMaxObserver, + granularity=QuantGranularity.PER_TENSOR, + verbose=verbose, + note="default with 16bit activation and 8bit weight", + ).add_node_target( + { + torch.ops.aten.conv2d.default, + }, + self.default_quant_dtype, + False, + act_observer=MinMaxObserver, + granularity=QuantGranularity.PER_CHANNEL, + note="Small model, so keep all conv2d at 16a8w per-channel instead of dropping to 4bit weights.", + ) + + +class Granite_3_3_2B_Instruct_HFQuantRecipe(LLMQuantRecipe): + """Ported from Granite_3_3_2B_InstructQuantRecipe (static_llm_quant_recipe.py:403). + + static_llama's ``layers\\..*\\.attention\\..*wv.*`` (the value projection) is + ``model.layers.N.self_attn.v_proj.conv`` in HF. + """ + + default_quant_dtype = QuantDtype.use_16a4w + + def __init__(self, verbose: bool = False): + super().__init__() + + self.recipe = ( + QuantRecipe( + self.default_quant_dtype, + False, + act_observer=MinMaxObserver, + granularity=QuantGranularity.PER_TENSOR, + verbose=verbose, + note="default with 16bit activation", + ) + .add_node_target( + { + torch.ops.aten.conv2d.default, + }, + QuantDtype.use_16a4w_block, + False, + act_observer=MinMaxObserver, + granularity=QuantGranularity.PER_BLOCK, + extra_kwargs={"block_size": (1, 64, 1, 1)}, + note="Annotate with 16a4w block quantization since these layers are not sensitive.", + ) + .add_regex( + { + r"model\.layers\..*\.self_attn\.v_proj\.conv", + }, + QuantDtype.use_16a8w, + False, + act_observer=MinMaxObserver, + granularity=QuantGranularity.PER_CHANNEL, + note="v_proj is sensitive and should be annotated with 16a8w.", + ) + ) + self.recipe.custom_quant_annotations.append(annotate_kv_8bit) + + +class DefaultQuantRecipe(LLMQuantRecipe): + """When quant recipe is not provided, this will be used""" + + default_quant_dtype = QuantDtype.use_16a8w + + def __init__(self, verbose: bool = False): + super().__init__() + self.recipe = QuantRecipe( + self.default_quant_dtype, + False, + act_observer=MinMaxObserver, + granularity=QuantGranularity.PER_TENSOR, + verbose=verbose, + note="default with 16bit activation and 8bit weight", + ).add_node_target( + { + torch.ops.aten.conv2d.default, + }, + QuantDtype.use_16a8w, + False, + act_observer=MinMaxObserver, + granularity=QuantGranularity.PER_CHANNEL, + note="Conservative fallback for models without a tuned recipe: 16a8w per-channel everywhere.", + ) + + +class Gemma3_1B_HFQuantRecipe(LLMQuantRecipe): + """Ported from Gemma3QuantRecipe (static_llm_quant_recipe.py:307).""" + + default_quant_dtype = QuantDtype.use_16a4w + masked_softmax = True + + def __init__(self, verbose: bool = False): + super().__init__() + + self.recipe = ( + QuantRecipe( + self.default_quant_dtype, + False, + act_observer=MinMaxObserver, + granularity=QuantGranularity.PER_TENSOR, + verbose=verbose, + note="default with 16bit activation", + ) + .add_node_target( + { + torch.ops.aten.conv2d.default, + }, + QuantDtype.use_16a4w_block, + False, + act_observer=MinMaxObserver, + granularity=QuantGranularity.PER_BLOCK, + extra_kwargs={"block_size": (1, 64, 1, 1)}, + note="Annotate with 16a4w block quantization since these layers are not sensitive.", + ) + .add_regex( + { + r"model\.layers\..*\.self_attn\.v_proj\.conv", + }, + QuantDtype.use_16a8w, + False, + act_observer=MinMaxObserver, + granularity=QuantGranularity.PER_CHANNEL, + note="v_proj is sensitive and should be annotated with 16a8w.", + ) + ) + self.recipe.custom_quant_annotations.append(annotate_kv_8bit) + + +HUGGING_FACE_QUANT_RECIPES = { + "google/gemma-3-1b-it": Gemma3_1B_HFQuantRecipe, + "NousResearch/Llama-3.2-1B": Llama3_2_1B_HFQuantRecipe, + "Qwen/Qwen2.5-0.5B": Qwen2_5_0_5B_HFQuantRecipe, + "Qwen/Qwen2.5-0.5B-Instruct": Qwen2_5_0_5B_HFQuantRecipe, + "Qwen/Qwen2.5-1.5B": Qwen2_5_1_5B_HFQuantRecipe, + "Qwen/Qwen2.5-1.5B-Instruct": Qwen2_5_1_5B_HFQuantRecipe, + "Qwen/Qwen3-0.6B": Qwen3_0_6B_HFQuantRecipe, + "HuggingFaceTB/SmolLM2-135M": Smollm2_HFQuantRecipe, + "HuggingFaceTB/SmolLM2-135M-Instruct": Smollm2_HFQuantRecipe, + "ibm-granite/granite-3.3-2b-instruct": Granite_3_3_2B_Instruct_HFQuantRecipe, +} + + +def get_quant_recipe(model_id: str, verbose: bool = True) -> LLMQuantRecipe: + """Return the tuned recipe for a Hugging Face model id, or the default one. + + model_id: The Hugging Face model id. + verbose: If True, the recipe logs the quant config tagged for each node. + """ + recipe_cls = HUGGING_FACE_QUANT_RECIPES.get(model_id, DefaultQuantRecipe) + if recipe_cls == DefaultQuantRecipe: + logging.warning( + f"{model_id} does not have customized quant recipe using default quant recipe." + ) + return recipe_cls(verbose) diff --git a/backends/qualcomm/runtime/QnnManager.cpp b/backends/qualcomm/runtime/QnnManager.cpp index 0f1169d26a8..b20a4fb4a84 100644 --- a/backends/qualcomm/runtime/QnnManager.cpp +++ b/backends/qualcomm/runtime/QnnManager.cpp @@ -14,6 +14,7 @@ #include #include #include +#include #include #include #include @@ -36,6 +37,35 @@ bool CompareExportedInput( return numA < numB; } +// The graph output position AOT encodes as `output_{id}_...`, or -1 for an +// output named without one (mutable buffers, older programs). +int ExportedOutputPosition(const std::string& name) { + const std::string prefix = "output_"; + if (name.rfind(prefix, 0) != 0) { + return -1; + } + const size_t end = name.find('_', prefix.size()); + if (end == std::string::npos || end == prefix.size() || + end + 1 == name.size()) { + return -1; + } + int position = -1; + const auto result = + std::from_chars(name.data() + prefix.size(), name.data() + end, position); + return result.ec == std::errc{} && result.ptr == name.data() + end && + position >= 0 + ? position + : -1; +} + +bool CompareExportedOutput( + const std::shared_ptr& a, + const std::shared_ptr& b) { + const int pa = ExportedOutputPosition(a->GetName()); + const int pb = ExportedOutputPosition(b->GetName()); + return pa >= 0 && (pb < 0 || pa < pb); +} + int ExtractMutableBufferNumber(const std::string& name) { std::string prefix = "mutbuf_"; size_t startPos = name.find(prefix); @@ -389,6 +419,14 @@ Error QnnManager::AllocateTensor(const std::string& graph_name) { } output_tensors_[graph_name].emplace_back(std::move(tensor_wrapper)); } + if (!options_->is_from_context_binary()) { + // Outputs with a position go in the order the program's outputs bind; + // the rest keep their place after them. + std::stable_sort( + output_tensors_[graph_name].begin(), + output_tensors_[graph_name].end(), + CompareExportedOutput); + } return Error::Ok; } @@ -406,6 +444,12 @@ Error QnnManager::AllocateTensor( CompareExportedInput); } output_tensors_[graph_name] = std::move(outputs); + if (!options_->is_from_context_binary()) { + std::stable_sort( + output_tensors_[graph_name].begin(), + output_tensors_[graph_name].end(), + CompareExportedOutput); + } return Error::Ok; } diff --git a/backends/qualcomm/tests/CMakeLists.txt b/backends/qualcomm/tests/CMakeLists.txt new file mode 100644 index 00000000000..df046000742 --- /dev/null +++ b/backends/qualcomm/tests/CMakeLists.txt @@ -0,0 +1,54 @@ +# 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. + +# Standalone Linux host tests: cmake -S backends/qualcomm/tests -B +cmake_minimum_required(VERSION 3.24) +project(qnn_prerequisite_tests LANGUAGES C CXX) +set(CMAKE_CXX_STANDARD 17) +set(CMAKE_CXX_STANDARD_REQUIRED ON) +get_filename_component( + EXECUTORCH_ROOT "${CMAKE_CURRENT_LIST_DIR}/../../.." ABSOLUTE +) +set(EXECUTORCH_BUILD_QNN ON CACHE BOOL "" FORCE) +set(EXECUTORCH_BUILD_EXTENSION_TENSOR ON CACHE BOOL "" FORCE) +set(EXECUTORCH_BUILD_TESTS OFF CACHE BOOL "" FORCE) +set(EXECUTORCH_BUILD_EXECUTOR_RUNNER OFF CACHE BOOL "" FORCE) +add_subdirectory("${EXECUTORCH_ROOT}" executorch) +add_subdirectory("${EXECUTORCH_ROOT}/third-party/googletest" googletest) +enable_testing() + +add_executable( + qnn_io_buffer_allocator_test test_io_buffer_allocator.cpp + ../llm/io_buffer_allocator.cpp +) +target_include_directories( + qnn_io_buffer_allocator_test PRIVATE + "${EXECUTORCH_ROOT}/.." + "${EXECUTORCH_ROOT}/runtime/core/portable_type/c10" + "${EXECUTORCH_ROOT}/third-party/flatbuffers/include" + "${CMAKE_BINARY_DIR}/schema/include" +) +target_link_libraries( + qnn_io_buffer_allocator_test PRIVATE executorch_core GTest::gtest_main +) +add_executable(qnn_output_order_test test_output_order.cpp) +target_include_directories( + qnn_output_order_test PRIVATE + "${EXECUTORCH_ROOT}/.." + "${EXECUTORCH_ROOT}/runtime/core/portable_type/c10" + "${EXECUTORCH_ROOT}/backends/qualcomm/runtime/pal/include" + "${EXECUTORCH_ROOT}/third-party/flatbuffers/include" + "${CMAKE_BINARY_DIR}/schema/include" + "$ENV{QNN_SDK_ROOT}/include/QNN" +) +target_link_libraries( + qnn_output_order_test PRIVATE qnn_manager qnn_backend_options + "$" + executorch_core GTest::gtest_main +) +include(GoogleTest) +gtest_discover_tests(qnn_io_buffer_allocator_test) +gtest_discover_tests(qnn_output_order_test) diff --git a/backends/qualcomm/tests/test_hf_static.py b/backends/qualcomm/tests/test_hf_static.py new file mode 100644 index 00000000000..d6bd1b49168 --- /dev/null +++ b/backends/qualcomm/tests/test_hf_static.py @@ -0,0 +1,629 @@ +# 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. + +"""Lowering Hugging Face's static-cache ABI must keep both the ABI and the +model. + +Tiny random decoders are exported the way Hugging Face's exporter does it (one +decode graph with a symbolic query length, the caches as inputs written in +place), specialized to a prompt chunk and a single token, rewritten, and driven +the way HF's generator drives the ABI. Greedy tokens must equal the model's +own, the inputs and outputs must be HF's, and in ExecuTorch the step must write +into the caller's cache buffers. +""" + +import builtins +import importlib.util +import itertools +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +import torch +import torch.utils._pytree as pytree + +try: + from transformers import ( + Gemma3ForCausalLM, + Gemma3TextConfig, + GenerationConfig, + LlamaConfig, + LlamaForCausalLM, + Qwen3Config, + Qwen3ForCausalLM, + ) + from transformers.exporters import DynamoConfig, DynamoExporter +except ImportError: # transformers without the exporters + DynamoExporter = None + +from executorch.backends.qualcomm.utils import hf_static +from executorch.extension.llm.io_contracts.hf_static import HfStaticDriver + +L = 32 # context length +N = 4 # prompt chunk +PROMPT = [5, 17, 42, 9, 77, 3, 64] # two chunks, the second padded +NEW_TOKENS = 8 +_SMALL = { + "hidden_size": 64, + "intermediate_size": 128, + "num_hidden_layers": 2, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "vocab_size": 128, + "max_position_embeddings": L, +} + + +def _llama(): + torch.manual_seed(0) + return LlamaForCausalLM(LlamaConfig(**_SMALL)).eval() + + +def _qwen3(): + torch.manual_seed(0) + return Qwen3ForCausalLM(Qwen3Config(head_dim=16, **_SMALL)).eval() + + +def _gemma3(): + torch.manual_seed(0) + # A window at least the context: HF's ABI has no wrap-around for sliding + # layers (their cache_position comes from a counter baked at 0). + return Gemma3ForCausalLM( + Gemma3TextConfig(head_dim=16, sliding_window=L, **_SMALL) + ).eval() + + +@unittest.skipIf(DynamoExporter is None, "needs transformers.exporters") +class TestHfStatic(unittest.TestCase): + make_model = staticmethod(_llama) + + def setUp(self): + torch.set_grad_enabled(False) + self.model = self.make_model() + self._export_hf_graph() + + def _export_hf_graph(self): + sample = torch.arange(4).reshape(1, 4) + 3 + self.hf_graph = ( + DynamoExporter() + .export_for_generation( + self.model, + {"input_ids": sample, "attention_mask": torch.ones_like(sample)}, + DynamoConfig(dynamic=True), + generation_config=GenerationConfig( + do_sample=False, cache_implementation="static", max_cache_len=L + ), + keep_all_logits=True, + ) + .components["decode"] + .artifact + ) + + def _programs(self, **options): + return { + q: hf_static.to_hf_static( + hf_static.specialize(self.hf_graph, q, L, config=self.model.config), + **options, + ) + for q in (1, N) + } + + def _expected(self): + return self.model.generate( + torch.tensor([PROMPT]), max_new_tokens=NEW_TOKENS, do_sample=False + )[0].tolist() + + def _generate(self, steps, layout, in_place=False): + driver = HfStaticDriver(steps, layout, in_place=in_place) + return driver.generate(PROMPT, NEW_TOKENS), driver + + def test_auto_attention_in_lowered_graphs(self): + for quantized in (False, True): + with self.subTest(quantized=quantized), patch.object( + hf_static, "_quantize", side_effect=lambda programs, *args: programs + ), patch.object( + hf_static.HfStaticPartitioner, "from_options", return_value=None + ), patch.object( + hf_static, "to_edge_transform_and_lower", return_value=None + ): + _, programs, _ = hf_static._lower( + self.hf_graph, + {"decode": 1, "prefill": N}, + L, + "auto", + quantized, + object() if quantized else None, + None, + ["prompt"], + "SM8750", + {}, + {}, + config=self.model.config, + ) + for method, program in programs.items(): + matmuls = [ + n + for n in program.graph.nodes + if n.target == torch.ops.aten.matmul.default + ] + expected = 4 if method == "decode" or quantized else 2 + self.assertEqual( + len(matmuls), expected * self.model.config.num_hidden_layers + ) + + def test_abi_is_hf(self): + programs = self._programs() + for q, program in programs.items(): + reference = hf_static.specialize( + self.hf_graph, q, L, config=self.model.config + ) + self.assertEqual( + [ + s.arg.name + for s in program.graph_signature.input_specs + if s.kind.name == "USER_INPUT" + ], + [ + s.arg.name + for s in reference.graph_signature.input_specs + if s.kind.name == "USER_INPUT" + ], + ) + for got, want in zip( + pytree.tree_leaves(program.example_inputs), + pytree.tree_leaves(reference.example_inputs), + ): + self.assertEqual((got.shape, got.dtype), (want.shape, want.dtype)) + self.assertFalse( + any("scaled_dot_product" in str(n.target) for n in program.graph.nodes) + ) + + def test_steps_match_hf_graph(self): + # Same inputs through the rewritten program and HF's own graph: a padded + # prompt chunk, then decode steps. Same logits (real rows), caches and + # counters, element by element. + programs = self._programs() + references = { + q: hf_static.specialize(self.hf_graph, q, L, config=self.model.config) + for q in programs + } + layout = hf_static.HfStaticLayout.of(programs[1]) + ours = HfStaticDriver( + {q: hf_static.flat_call(p) for q, p in programs.items()}, layout + ) + theirs = HfStaticDriver( + {q: hf_static.flat_call(p) for q, p in references.items()}, layout + ) + steps = [(PROMPT[:N], 0, N), (PROMPT[N:], N, N)] + steps += [([t], len(PROMPT) + i, 1) for i, t in enumerate([1, 2, 3])] + for tokens, start, q in steps: + torch.testing.assert_close( + ours.forward(tokens, start, q), theirs.forward(tokens, start, q) + ) + for a, b in zip(ours.state, theirs.state): + if a.dim() == 0: + torch.testing.assert_close(a, b) + else: + # Padding slots are not valid cache state under this contract. + filled = start + len(tokens) + torch.testing.assert_close(a[:, :, :filled], b[:, :, :filled]) + + def test_greedy_matches_model(self): + programs = self._programs() + layout = hf_static.HfStaticLayout.of(programs[1]) + tokens, _ = self._generate( + {q: hf_static.flat_call(p) for q, p in programs.items()}, layout + ) + self.assertEqual(tokens, self._expected()) + + def test_greedy_matches_model_masked_softmax(self): + programs = self._programs(masked_softmax=True) + layout = hf_static.HfStaticLayout.of(programs[1]) + tokens, _ = self._generate( + {q: hf_static.flat_call(p) for q, p in programs.items()}, layout + ) + self.assertEqual(tokens, self._expected()) + + def test_greedy_matches_model_attention_forms(self): + for attention, masked_softmax in itertools.product( + ("split", "concat"), (False, True) + ): + programs = self._programs( + attention=attention, masked_softmax=masked_softmax + ) + layout = hf_static.HfStaticLayout.of(programs[1]) + tokens, _ = self._generate( + {q: hf_static.flat_call(p) for q, p in programs.items()}, layout + ) + self.assertEqual(tokens, self._expected()) + + def test_executorch_writes_inputs_in_place(self): + from executorch.exir import EdgeCompileConfig, ExecutorchBackendConfig, to_edge + from executorch.exir.passes import MemoryPlanningPass + + programs = self._programs() + edge = to_edge( + {f"q{q}": p for q, p in programs.items()}, + compile_config=EdgeCompileConfig(_check_ir_validity=False), + ) + layout = hf_static.HfStaticLayout.of(programs[1]) + et = edge.to_executorch( + ExecutorchBackendConfig( + run_reinplace_pass=True, + memory_planning_pass=MemoryPlanningPass( + alloc_graph_input=False, alloc_graph_output=False + ), + ) + ) + # Each cache write and counter update is an in-place kernel on the + # method's own input value, so it lands in the buffer the caller binds. + for plan in et.executorch_program.execution_plan: + ops = [o.name for o in plan.operators] + inputs = set(plan.inputs) + writes = [ + (ops[i.instr_args.op_index], i.instr_args.args[0]) + for chain in plan.chains + for i in chain.instructions + if hasattr(i.instr_args, "op_index") + and ops[i.instr_args.op_index] in ("aten::index_put_", "aten::copy_") + ] + windows = hf_static.sliding_windows(self.model.config) + self.assertEqual( + len(writes), + 3 * layout.n_layers + 2 * sum(w is not None for w in windows), + ) + self.assertTrue(all(value in inputs for _, value in writes)) + + def test_qnn_partitions(self): + try: + from executorch.exir import EdgeCompileConfig, to_edge_transform_and_lower + except ImportError: + self.skipTest("needs ExecuTorch") + try: + partitioner = hf_static.HfStaticPartitioner.from_options( + soc_model="SM8750", use_fp16=True, shared_buffer=True + ) + except Exception as error: # QNN SDK not set up + self.skipTest(f"needs the QNN SDK: {error}") + program = self._programs()[N] + edge = to_edge_transform_and_lower( + program, + partitioner=[partitioner], + compile_config=EdgeCompileConfig(_check_ir_validity=False), + ) + calls = [ + n + for n in edge.exported_program().graph.nodes + if n.op == "call_function" and "getitem" not in str(n.target) + ] + delegates = [n for n in calls if "executorch_call_delegate" in str(n.target)] + self.assertEqual(len(delegates), 1) + on_cpu = {str(n.target) for n in calls if n not in delegates} + self.assertTrue(any("index_put" in t for t in on_cpu)) + self.assertTrue(any("logical_and" in t for t in on_cpu)) + + def test_export_hf_static(self): + try: + hf_static.HfStaticPartitioner.from_options( + soc_model="SM8750", use_fp16=True + ) + except Exception as error: # QNN SDK not set up + self.skipTest(f"needs the QNN SDK: {error}") + + class Tokenizer: + def __call__(self, text, return_tensors=None): + return {"input_ids": torch.tensor([PROMPT])} + + # Parameters require grad, as from_pretrained loads them. + with torch.enable_grad(): + program = hf_static.export_hf_static( + self.make_model(), Tokenizer(), "SM8750", max_seq_len=L, prefill_len=N + ) + plans = {p.name for p in program.executorch_program.execution_plan} + self.assertLessEqual( + {hf_static.PREFILL_METHOD, hf_static.DECODE_METHOD, "get_io_contract"}, + plans, + ) + + def _quantizer(self, custom_annotations=()): + try: + from executorch.backends.qualcomm.export_utils import ( + get_backend_type, + make_quantizer, + ) + from executorch.backends.qualcomm.quantizer.quantizer import QuantDtype + from torchao.quantization.pt2e import MinMaxObserver + except Exception as error: # QNN SDK not set up + self.skipTest(f"needs the QNN quantizer: {error}") + # Min/max activation observers, as the LLM recipes use. + return make_quantizer( + quant_dtype=QuantDtype.use_16a8w, + act_observer=MinMaxObserver, + custom_annotations=custom_annotations, + backend=get_backend_type("htp"), + soc_model="SM8750", + ) + + def _quantize(self, programs, seq_mse_candidates=0, quantizer=None): + class Tokenizer: + def __call__(self, text): + return {"input_ids": list(PROMPT)} + + programs = { + hf_static.DECODE_METHOD: programs[1], + hf_static.PREFILL_METHOD: programs[N], + } + lengths = {hf_static.DECODE_METHOD: 1, hf_static.PREFILL_METHOD: N} + layout = hf_static.HfStaticLayout.of(programs[hf_static.DECODE_METHOD]) + quantized = hf_static._quantize( + programs, + lengths, + layout, + quantizer or self._quantizer(), + Tokenizer(), + ["hello world"], + seq_mse_candidates, + ) + return programs, quantized, layout + + @staticmethod + def _steps(programs): + return { + 1: hf_static.flat_call(programs[hf_static.DECODE_METHOD]), + N: hf_static.flat_call(programs[hf_static.PREFILL_METHOD]), + } + + def _quantized_fidelity(self, programs, seq_mse_candidates=0): + """Cosine similarity of the quantized programs' logits to the float + ones, over a prompt chunk and two decode steps.""" + programs, quantized, layout = self._quantize(programs, seq_mse_candidates) + drivers = [ + HfStaticDriver(self._steps(p), layout) for p in (programs, quantized) + ] + fidelity = [] + for tokens, start, q in [(PROMPT[:N], 0, N), ([1], N, 1), ([2], N + 1, 1)]: + ref, got = (d.forward(tokens, start, q) for d in drivers) + fidelity.append( + torch.nn.functional.cosine_similarity(ref, got, dim=0).item() + ) + return fidelity + + def test_kv_8bit_annotation_keeps_float_cache_inputs(self): + try: + from executorch.backends.qualcomm.export_utils import ( + get_backend_type, + make_quantizer, + ) + from executorch.backends.qualcomm.quantizer.custom_annotation import ( + annotate_kv_8bit, + ) + from executorch.backends.qualcomm.quantizer.quantizer import QuantDtype + except Exception as error: # QNN SDK not set up + self.skipTest(f"needs the QNN quantizer: {error}") + from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e + + quantizer = make_quantizer( + quant_dtype=QuantDtype.use_16a8w, + custom_annotations=(annotate_kv_8bit,), + backend=get_backend_type("htp"), + soc_model="SM8750", + ) + for program in self._programs(masked_softmax=True).values(): + prepared = prepare_pt2e(program.module(), quantizer) + args, kwargs = program.example_inputs + prepared(*args, **kwargs) + graph = convert_pt2e(prepared).graph + caches = [ + n + for n in graph.nodes + if n.op == "placeholder" + and n.name.startswith("past_key_values_") + and n.meta["val"].dim() == 4 + ] + self.assertTrue(caches) + for cache in caches: + self.assertEqual(cache.meta["val"].dtype, torch.float32) + self.assertIn( + torch.uint8, + {u.args[-1] for u in cache.users if "quantize" in str(u.target)}, + cache.name, + ) + + def test_quantized_tracks_hf_graph_quantized(self): + # The rewrite costs no accuracy under quantization: on average over the + # steps, as close to float as HF's own graph quantized the same way. + hf = self._quantized_fidelity( + { + q: hf_static.specialize(self.hf_graph, q, L, config=self.model.config) + for q in (1, N) + } + ) + ours = self._quantized_fidelity(self._programs(masked_softmax=True)) + self.assertGreater(sum(ours) / len(ours), sum(hf) / len(hf) - 0.01) + + def test_quantize_with_seq_mse(self): + fidelity = self._quantized_fidelity(self._programs(masked_softmax=True), 4) + self.assertGreater(min(fidelity), 0.9) + + +@unittest.skipIf(DynamoExporter is None, "needs transformers.exporters") +class TestHfStaticSlidingWindow(unittest.TestCase): + """Gemma3 with a sliding window shorter than the context: generation past + the window, prompt chunks that roll it (the last one padded).""" + + WINDOW = 6 + PROMPT = [5, 17, 42, 9, 77, 3, 64, 33, 21, 8, 99] # three chunks + NEW_TOKENS = 20 + + def setUp(self): + torch.set_grad_enabled(False) + torch.manual_seed(0) + config = dict(_SMALL, num_hidden_layers=6) # five sliding layers, one full + self.model = Gemma3ForCausalLM( + Gemma3TextConfig(head_dim=16, sliding_window=self.WINDOW, **config) + ).eval() + sample = torch.arange(4).reshape(1, 4) + 3 + self.hf_graph = ( + DynamoExporter() + .export_for_generation( + self.model, + {"input_ids": sample, "attention_mask": torch.ones_like(sample)}, + DynamoConfig(dynamic=True), + generation_config=GenerationConfig( + do_sample=False, cache_implementation="static", max_cache_len=L + ), + keep_all_logits=True, + ) + .components["decode"] + .artifact + ) + + def _programs(self, **options): + return { + q: hf_static.to_hf_static( + hf_static.specialize(self.hf_graph, q, L, config=self.model.config), + **options, + ) + for q in (1, N) + } + + def _assert_greedy_matches_model(self, programs, steps=None, in_place=False): + layout = hf_static.HfStaticLayout.of(programs[1]) + steps = steps or {q: hf_static.flat_call(p) for q, p in programs.items()} + tokens = HfStaticDriver(steps, layout, in_place).generate( + self.PROMPT, self.NEW_TOKENS + ) + expected = self.model.generate( + torch.tensor([self.PROMPT]), + max_new_tokens=self.NEW_TOKENS, + do_sample=False, + )[0].tolist() + # HF stops at the end-of-sequence token; past the window either way. + self.assertGreater(len(expected), len(self.PROMPT) + 2) + self.assertGreater(len(expected), self.WINDOW) + self.assertEqual(tokens[: len(expected)], expected) + + def test_window_caches_keep_hf_shapes(self): + reference = hf_static.specialize(self.hf_graph, 1, L, config=self.model.config) + program = self._programs()[1] + for got, want in zip( + pytree.tree_leaves(program.example_inputs), + pytree.tree_leaves(reference.example_inputs), + ): + self.assertEqual((got.shape, got.dtype), (want.shape, want.dtype)) + layout = hf_static.HfStaticLayout.of(program) + self.assertEqual(sorted({s[2] for s in layout.cache_shapes}), [self.WINDOW, L]) + + def test_window_greedy_matches_model(self): + for attention, masked_softmax in itertools.product( + ("split", "concat"), (False, True) + ): + with self.subTest(attention=attention, masked_softmax=masked_softmax): + self._assert_greedy_matches_model( + self._programs(attention=attention, masked_softmax=masked_softmax) + ) + + def test_window_writes_in_place_in_executorch(self): + # The shift and the new rows are in-place kernels on the method's own + # cache inputs; run through ExecuTorch with the caller keeping its + # buffers, generation still matches HF. + from executorch.exir import EdgeCompileConfig, ExecutorchBackendConfig, to_edge + from executorch.exir.passes import MemoryPlanningPass + + try: + from executorch.extension.pybindings.portable_lib import ( + _load_for_executorch_from_buffer, + ) + except ImportError: + self.skipTest("needs ExecuTorch's runtime") + programs = self._programs(masked_softmax=True) + layout = hf_static.HfStaticLayout.of(programs[1]) + edge = to_edge( + {f"q{q}": p for q, p in programs.items()}, + compile_config=EdgeCompileConfig(_check_ir_validity=False), + ) + et = edge.to_executorch( + ExecutorchBackendConfig( + run_reinplace_pass=True, + memory_planning_pass=MemoryPlanningPass( + alloc_graph_input=False, alloc_graph_output=False + ), + ) + ) + module = _load_for_executorch_from_buffer(et.buffer) + mutated = 3 * layout.n_layers + + def step(q): + def run(*inputs): + # The mutated inputs come back first, then the logits. + outputs = module.run_method(f"q{q}", inputs) + return [outputs[mutated], *outputs[:mutated]] + + return run + + self._assert_greedy_matches_model( + programs, {q: step(q) for q in programs}, in_place=True + ) + + +class TestHfStaticQwen3(TestHfStatic): + make_model = staticmethod(_qwen3) + + +class TestHfStaticGemma3(TestHfStatic): + make_model = staticmethod(_gemma3) + + +class TestHfStaticCpuTags(unittest.TestCase): + def test_import_without_exporters_or_pt2e(self): + original_import = builtins.__import__ + + def import_without_export(name, *args, **kwargs): + if name in ( + "transformers.exporters", + "torchao.quantization.pt2e.quantize_pt2e", + ): + raise ImportError("dependency deliberately unavailable") + return original_import(name, *args, **kwargs) + + spec = importlib.util.spec_from_file_location( + "hf_static_without_export", hf_static.__file__ + ) + module = importlib.util.module_from_spec(spec) + with patch("builtins.__import__", side_effect=import_without_export): + spec.loader.exec_module(module) + self.assertIsNone(module.DynamoExporter) + self.assertIsNone(module.prepare_pt2e) + + def test_export_prefill_window_guard_precedes_tracing(self): + model = SimpleNamespace( + config=SimpleNamespace(num_hidden_layers=2, sliding_window=6) + ) + with self.assertRaisesRegex( + ValueError, "prefill_len.*smallest sliding cache.*6" + ): + hf_static.export_hf_static( + model, None, "SM8750", max_seq_len=32, prefill_len=8 + ) + + def test_cpu_selection_uses_tags_not_dtype_size_or_name(self): + graph = torch.fx.Graph() + tokens = graph.placeholder("input_ids") + tokens.meta["val"] = torch.zeros(1, 4, dtype=torch.int64) + untagged = graph.call_function( + torch.ops.aten._to_copy.default, (tokens,), {"dtype": torch.int64} + ) + untagged.meta["val"] = torch.zeros(4, dtype=torch.int64) + tagged = graph.call_function(torch.ops.aten.sin.default, (tokens,)) + tagged.meta = {"val": torch.zeros(128), "custom": {hf_static.CPU_META: True}} + graph.output((untagged, tagged)) + names = hf_static._cpu_nodes(SimpleNamespace(graph=graph)) + self.assertIn(tagged.name, names) + self.assertNotIn(untagged.name, names) + + +if __name__ == "__main__": + unittest.main() diff --git a/backends/qualcomm/tests/test_io_buffer_allocator.cpp b/backends/qualcomm/tests/test_io_buffer_allocator.cpp new file mode 100644 index 00000000000..024fe19c4e3 --- /dev/null +++ b/backends/qualcomm/tests/test_io_buffer_allocator.cpp @@ -0,0 +1,62 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include +#include + +namespace { +void* allocated_block; +size_t requested_bytes; +size_t requested_alignment; +void* freed_block; +void* bound_tensor; +void* bound_block; +} // namespace + +void* QnnExecuTorchAllocCustomMem(size_t bytes, size_t alignment) { + requested_bytes = bytes; + requested_alignment = alignment; + return allocated_block; +} + +void QnnExecuTorchFreeCustomMem(void* block) { + freed_block = block; +} + +void QnnExecuTorchAddCustomMemTensorAddr(void* tensor, void* block) { + bound_tensor = tensor; + bound_block = block; +} + +TEST(QnnIOBufferAllocatorTest, RegistersAndForwardsAllocationBindingAndFree) { + auto& factory = executorch::extension::llm::io_buffer_allocator_factory(); + ASSERT_TRUE(factory); + auto allocator = factory(); + ASSERT_NE(allocator, nullptr); + alignas(64) char memory[128]; + allocated_block = memory; + EXPECT_EQ(allocator->allocate(sizeof(memory), 64), memory); + EXPECT_EQ(requested_bytes, sizeof(memory)); + EXPECT_EQ(requested_alignment, 64); + allocator->bind(memory, memory + 32); + EXPECT_EQ(bound_tensor, memory + 32); + EXPECT_EQ(bound_block, memory); + allocator->deallocate(memory); + EXPECT_EQ(freed_block, memory); +} + +TEST(QnnIOBufferAllocatorTest, PreservesFailedAllocationForHeapFallback) { + auto& factory = executorch::extension::llm::io_buffer_allocator_factory(); + ASSERT_TRUE(factory); + allocated_block = nullptr; + auto allocator = factory(); + EXPECT_EQ(allocator->allocate(1024, 64), nullptr); + EXPECT_EQ(requested_bytes, 1024); + EXPECT_EQ(requested_alignment, 64); +} diff --git a/backends/qualcomm/tests/test_kv_cache_io.py b/backends/qualcomm/tests/test_kv_cache_io.py new file mode 100644 index 00000000000..29bd3ef3629 --- /dev/null +++ b/backends/qualcomm/tests/test_kv_cache_io.py @@ -0,0 +1,300 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""The KV-cache-as-I/O rewrite must reproduce the model it rewrites. + +A tiny random Llama is exported the way Hugging Face exporters do it (static +cache as inputs), rewritten, and driven the way qnn_llama_runner drives it +(chunked prompt step, then decode steps writing the new slice into the past +buffers under the smart mask). Its greedy tokens must equal the model's own. +""" + +import unittest + +import torch + +try: + from transformers import ( + GenerationConfig, + LlamaConfig, + LlamaForCausalLM, + Qwen3Config, + Qwen3ForCausalLM, + ) + from transformers.exporters import DynamoConfig, DynamoExporter +except ImportError: # transformers without the exporters + DynamoExporter = None + +from executorch.backends.qualcomm.utils import kv_cache_io + +L = 32 # context length + + +def _model(): + torch.manual_seed(0) + config = LlamaConfig( + hidden_size=64, + intermediate_size=128, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + vocab_size=128, + max_position_embeddings=L, + ) + return LlamaForCausalLM(config).eval() + + +def _qwen3_model(): + torch.manual_seed(0) + config = Qwen3Config( + hidden_size=64, + intermediate_size=128, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=16, + vocab_size=128, + max_position_embeddings=L, + ) + return Qwen3ForCausalLM(config).eval() + + +@unittest.skipIf(DynamoExporter is None, "needs transformers.exporters") +class TestKvCacheIo(unittest.TestCase): + make_model = staticmethod(_model) + + def setUp(self): + self.model = self.make_model() + sample = torch.arange(4).reshape(1, 4) + 3 + self.artifacts = artifacts = DynamoExporter().export_for_generation( + self.model, + {"input_ids": sample, "attention_mask": torch.ones_like(sample)}, + DynamoConfig(), + generation_config=GenerationConfig( + do_sample=False, cache_implementation="static", max_cache_len=L + ), + ) + self.prefill = kv_cache_io.to_kv_cache_io( + artifacts.components["prefill"].artifact + ) + self.decode = kv_cache_io.to_kv_cache_io( + artifacts.components["decode"].artifact + ) + + def test_contract(self): + self.assertEqual(self.decode.method_name, "kv_forward") + self.assertEqual(self.prefill.method_name, "prefill_forward") + self.assertEqual(self.decode.k_shapes[0], torch.Size([1, 2, 16, L - 1])) + self.assertEqual(self.decode.v_shapes[0], torch.Size([1, 2, L - 1, 16])) + self.assertEqual(tuple(self.prefill.logits_shape), (1, 4, 128)) + self.assertEqual( + self.decode.constant_methods()["get_io_contract"], "kvcache_io" + ) + targets = {str(n.target) for n in self.decode.module.graph.nodes} + self.assertFalse( + any("index_copy" in t or "scaled_dot_product" in t for t in targets) + ) + # RoPE runs in halves, without building rotate_half's rotated copy. + self.assertNotIn("aten.neg.default", targets) + + def test_weight_arithmetic_is_folded(self): + # Gemma-style RMSNorm scale: (1 + weight) becomes one constant. + class Norm(torch.nn.Module): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.randn(8)) + + def forward(self, x): + norm = x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + 1e-6) + return norm * (1.0 + self.weight.float()) + + gm = torch.export.export(Norm(), (torch.randn(2, 8),)).module() + x = torch.randn(2, 8) + expected = gm(x) + kv_cache_io.fold_weight_arithmetic(gm) + gm.recompile() + last = next(n for n in gm.graph.nodes if n.op == "output").args[0][0] + self.assertTrue(any(a.op == "get_attr" for a in last.args)) + torch.testing.assert_close(gm(x), expected) + # The rewritten step re-exports cleanly. + self.decode.export() + + def test_greedy_matches_model(self): + self._check_greedy(self.decode, self.prefill, -65535.0) + + def test_greedy_matches_model_masked_softmax(self): + artifacts = self.artifacts.components + self._check_greedy( + kv_cache_io.to_kv_cache_io( + artifacts["decode"].artifact, masked_softmax=True + ), + kv_cache_io.to_kv_cache_io( + artifacts["prefill"].artifact, masked_softmax=True + ), + kv_cache_io.CALIBRATION_MASK_VALUE, + ) + + def test_greedy_matches_model_r3(self): + artifacts = self.artifacts.components + self._check_greedy( + kv_cache_io.to_kv_cache_io(artifacts["decode"].artifact, r3=True), + kv_cache_io.to_kv_cache_io(artifacts["prefill"].artifact, r3=True), + -65535.0, + ) + + def _check_greedy(self, dec, pre, masked): + layers = dec.num_layers + prompt = [5, 17, 42, 9, 77, 3, 64] # two prompt chunks of 4, the second padded + past_k = [torch.zeros(s) for s in dec.k_shapes] + past_v = [torch.zeros(s) for s in dec.v_shapes] + ar, past = pre.ar_len, L - pre.ar_len + with torch.no_grad(): + start = 0 + while start < len(prompt): + chunk = prompt[start : start + ar] + n = len(chunk) + mask = torch.full((1, 1, ar, L), masked) + for r in range(ar): + mask[..., r, :start] = 0 + mask[..., r, past : past + min(r, n - 1) + 1] = 0 + out = pre.module( + torch.tensor([chunk + [chunk[-1]] * (ar - n)], dtype=torch.int32), + mask, + torch.tensor([[start + i for i in range(ar)]], dtype=torch.int32), + *[k[..., :past] for k in past_k], + *[v[..., :past, :] for v in past_v], + ) + for i in range(layers): + past_k[i][..., start : start + n] = out[1 + i][..., :n] + past_v[i][..., start : start + n, :] = out[1 + layers + i][ + ..., :n, : + ] + last = out[0][0, n - 1] + start += n + tokens = prompt + [int(last.argmax())] + for _ in range(8): + pos = len(tokens) - 1 + mask = torch.full((1, 1, 1, L), masked) + mask[..., :pos] = 0 + mask[..., L - 1 :] = 0 + out = dec.module( + torch.tensor([[tokens[-1]]], dtype=torch.int32), + mask, + torch.tensor([[pos]], dtype=torch.int32), + *past_k, + *past_v, + ) + for i in range(layers): + past_k[i][..., pos] = out[1 + i][..., 0] + past_v[i][..., pos, :] = out[1 + layers + i][..., 0, :] + tokens.append(int(out[0][0, -1].argmax())) + expected = self.model.generate( + torch.tensor([prompt]), max_new_tokens=9, do_sample=False + )[0].tolist() + self.assertEqual(tokens, expected) + + +class TestKvCacheIoQkNorm(TestKvCacheIo): + """Qwen3: per-head query and key norms.""" + + make_model = staticmethod(_qwen3_model) + + def test_qk_norm_after_head_transpose(self): + for step in (self.decode, self.prefill): + means = [ + n + for n in step.export().graph.nodes + if n.target == torch.ops.aten.mean.dim and n.meta["val"].dim() == 4 + ] + self.assertEqual(len(means), 4) # q and k norms of 2 layers + for mean in means: + # Normalized as [B, H, S, D], the layout the attention reads. + _, heads, steps, _ = mean.meta["val"].shape + self.assertIn(heads, (2, 4)) + self.assertEqual(steps, step.ar_len) + + +@unittest.skipIf(DynamoExporter is None, "needs transformers.exporters") +class TestKvCacheIoSeq2Seq(unittest.TestCase): + """An encoder-decoder (tiny random Whisper): self-attention caches as I/O, + cross-attention keys and values as fixed inputs, positions from the + caches' counters.""" + + def test_greedy_matches_model(self): + from executorch.backends.qualcomm.utils.hf_export import _CrossAttentionInputs + from transformers import WhisperConfig, WhisperForConditionalGeneration + + torch.manual_seed(0) + config = WhisperConfig( + vocab_size=64, + num_mel_bins=8, + d_model=32, + encoder_layers=2, + decoder_layers=2, + encoder_attention_heads=4, + decoder_attention_heads=4, + encoder_ffn_dim=64, + decoder_ffn_dim=64, + max_source_positions=16, + max_target_positions=L, + decoder_start_token_id=1, + eos_token_id=2, + pad_token_id=0, + bos_token_id=1, + suppress_tokens=[], + begin_suppress_tokens=[], + ) + model = WhisperForConditionalGeneration(config).eval() + features = torch.randn(1, 8, 32) + generation = GenerationConfig( + do_sample=False, + cache_implementation="static", + max_cache_len=L, + max_new_tokens=8, + decoder_start_token_id=1, + eos_token_id=2, + ) + artifacts = DynamoExporter().export_for_generation( + model, + {"input_features": features}, + DynamoConfig(), + generation_config=generation, + ) + step = kv_cache_io.to_kv_cache_io(artifacts.components["decode"].artifact) + self.assertEqual(len(step.cross_k_shapes), 2) + self.assertEqual(step.constant_methods()["get_n_cross_layers"], 2) + with torch.no_grad(): + cross = _CrossAttentionInputs(model, 1)(features) + past_k = [torch.zeros(s) for s in step.k_shapes] + past_v = [torch.zeros(s) for s in step.v_shapes] + tokens = [1] + for pos in range(8): + mask = torch.full((1, 1, 1, L), -65535.0) + mask[..., :pos] = 0 + mask[..., L - 1 :] = 0 + out = step.module( + torch.tensor([[tokens[-1]]], dtype=torch.int32), + mask, + torch.tensor([[pos]], dtype=torch.int32), + *past_k, + *past_v, + *cross, + ) + for i in range(step.num_layers): + past_k[i][..., pos] = out[1 + i][..., 0] + past_v[i][..., pos, :] = out[1 + step.num_layers + i][..., 0, :] + tokens.append(int(out[0][0, -1].argmax())) + expected = [1] + for _ in range(8): + logits = model( + input_features=features, decoder_input_ids=torch.tensor([expected]) + ).logits + expected.append(int(logits[0, -1].argmax())) + self.assertEqual(tokens, expected) + + +if __name__ == "__main__": + unittest.main() diff --git a/backends/qualcomm/tests/test_output_order.cpp b/backends/qualcomm/tests/test_output_order.cpp new file mode 100644 index 00000000000..0f95b6649ac --- /dev/null +++ b/backends/qualcomm/tests/test_output_order.cpp @@ -0,0 +1,157 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include + +#include +#include +#include + +namespace executorch::backends::qnn { +int ExportedOutputPosition(const std::string& name); + +namespace { + +std::shared_ptr output( + const std::string& name, + float* data = nullptr) { + const uint32_t dims[] = {2}; + auto tensor = CreateTensorWrapper( + name, + QNN_TENSOR_TYPE_APP_READ, + QNN_DATATYPE_FLOAT_32, + std::make_unique(), + 1, + dims, + nullptr); + if (data != nullptr) { + tensor->FillDataBuffer(data); + } + return tensor; +} + +class QnnOutputOrderTest : public ::testing::Test { + protected: + void SetUp() override { + auto backend = qnn_delegate::CreateQnnExecuTorchBackendOptions( + builder_, qnn_delegate::QnnExecuTorchBackendType::kHtpBackend); + qnn_delegate::QnnExecuTorchOptionsBuilder options(builder_); + options.add_backend_options(backend); + options.add_is_from_context_binary(from_context_binary()); + builder_.Finish(options.Finish()); + manager_ = std::make_unique( + qnn_delegate::GetQnnExecuTorchOptions(builder_.GetBufferPointer()), + QnnExecuTorchContextBinary{nullptr, 0}); + } + + virtual bool from_context_binary() const { + return false; + } + + std::vector> allocate( + std::vector> tensors) { + std::vector> inputs; + EXPECT_EQ( + manager_->AllocateTensor("forward", inputs, tensors), + executorch::runtime::Error::Ok); + return manager_->GetGraphOutputs("forward"); + } + + flatbuffers::FlatBufferBuilder builder_; + std::unique_ptr manager_; +}; + +TEST_F(QnnOutputOrderTest, ReorderedSameShapedOutputsBindDistinctValues) { + std::array key = {11, 12}; + std::array value = {21, 22}; + auto k = output("output_0_key@0", key.data()); + auto v = output("output_1_value@0", value.data()); + auto tensors = allocate({v, k}); + ASSERT_EQ(tensors.size(), 2); + EXPECT_EQ(tensors[0], k); + EXPECT_EQ(tensors[1], v); + EXPECT_EQ(tensors[0]->GetStaticTensorData(), key.data()); + EXPECT_EQ(tensors[1]->GetStaticTensorData(), value.data()); + EXPECT_FLOAT_EQ( + static_cast(tensors[0]->GetStaticTensorData())[0], 11); + EXPECT_FLOAT_EQ( + static_cast(tensors[1]->GetStaticTensorData())[0], 21); +} + +TEST_F(QnnOutputOrderTest, NumericOrderAndEqualPositionsAreStable) { + auto zero = output("output_0_zero@0"); + auto two = output("output_2_two@0"); + auto duplicate = output("output_2_duplicate@0"); + auto ten = output("output_10_ten@0"); + EXPECT_EQ( + allocate({ten, two, duplicate, zero}), + (std::vector>{zero, two, duplicate, ten})); +} + +TEST_F( + QnnOutputOrderTest, + LegacyMutableAndDumpOutputsKeepRelativeOrderAndData) { + float cache[2] = {4, 5}; + auto legacy = output("output_old@0"); + auto mutated = output("output_mutbuf_0_cache@0", cache); + auto dump = output("intermediate@0"); + EXPECT_EQ( + allocate({legacy, mutated, dump}), + (std::vector>{legacy, mutated, dump})); + auto numbered = output("output_0_new@0"); + EXPECT_EQ( + allocate({legacy, numbered, mutated, dump}), + (std::vector>{ + numbered, legacy, mutated, dump})); + EXPECT_EQ(mutated->GetStaticTensorData(), cache); + EXPECT_EQ(mutated->GetName(), "output_mutbuf_0_cache@0"); +} + +class QnnContextOutputOrderTest : public QnnOutputOrderTest { + protected: + bool from_context_binary() const override { + return true; + } +}; + +TEST_F(QnnContextOutputOrderTest, ContextBinaryOrderIsNotChanged) { + auto one = output("output_1_first@0"); + auto zero = output("output_0_second@0"); + EXPECT_EQ( + allocate({one, zero}), + (std::vector>{one, zero})); +} + +TEST( + QnnOutputPositionTest, + RejectsLegacyMalformedAndOverflowNamesWithoutThrowing) { + for (const std::string name : + {"output_old@0", + "output_mutbuf_0_cache@0", + "intermediate@0", + "output_", + "output_2", + "output_2_", + "output_2old@0", + "output_-1_old@0", + "output_+1_old@0", + "output_ 1_old@0", + "output_2147483648_old@0", + "output_999999999999999999999_old@0"}) { + EXPECT_NO_THROW(EXPECT_EQ(ExportedOutputPosition(name), -1)) << name; + } + EXPECT_EQ(ExportedOutputPosition("output_0_add@0"), 0); + EXPECT_EQ(ExportedOutputPosition("output_10_getitem@2"), 10); + EXPECT_EQ( + ExportedOutputPosition("output_2147483647_last@0"), + std::numeric_limits::max()); +} + +} // namespace +} // namespace executorch::backends::qnn diff --git a/backends/qualcomm/tests/test_output_order.py b/backends/qualcomm/tests/test_output_order.py new file mode 100644 index 00000000000..da8d5a13af2 --- /dev/null +++ b/backends/qualcomm/tests/test_output_order.py @@ -0,0 +1,122 @@ +# 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. + +import operator +from collections import defaultdict +from types import SimpleNamespace + +import executorch.backends.qualcomm.python.PyQnnManagerAdaptor as PyQnnManager + +import pytest +import torch +from executorch.backends.qualcomm.builders.node_visitor import NodeVisitor +from executorch.backends.qualcomm.utils.constants import QCOM_TENSOR_NAME + + +def _visitor(graph, *, buffers=None, mutations=None, validation=False, dump=False): + program = SimpleNamespace( + graph_signature=SimpleNamespace( + inputs_to_buffers=buffers or {}, + buffers_to_mutate=mutations or {}, + inputs_to_parameters={}, + inputs_to_lifted_tensor_constants={}, + ) + ) + return NodeVisitor( + {node: index for index, node in enumerate(graph.nodes)}, + program, + enable_tensor_dump=dump, + is_qnn_partitioner=not validation, + ) + + +def test_reordered_same_shaped_outputs_keep_program_positions(): + graph = torch.fx.Graph() + x = graph.placeholder("x") + add = graph.call_function(torch.add, (x, 7)) + neg = graph.call_function(torch.neg, (x,)) + graph.output((neg, add)) + module = torch.fx.GraphModule({}, graph) + sample = torch.tensor([1.0, 2.0]) + expected = module(sample) + visitor = _visitor(graph, dump=True) + wrappers = defaultdict(dict) + built = [ + visitor.define_tensor( + node, + node, + value, + PyQnnManager.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE, + wrappers, + ) + for node, value in ((add, sample + 7), (neg, -sample)) + ] + built = [PyQnnManager.PyQnnTensorWrapper(wrapper) for wrapper in built] + assert [wrapper.GetName() for wrapper in built] == [ + f"output_1_{add.name}@0", + f"output_0_{neg.name}@0", + ] + assert add.meta[QCOM_TENSOR_NAME][0] == built[0].GetName() + assert neg.meta[QCOM_TENSOR_NAME][0] == built[1].GetName() + assert expected[0].shape == expected[1].shape + assert not torch.equal(expected[0], expected[1]) + + +@pytest.mark.parametrize("escaping", [(2, 0), (1,), (0, 2, 1)]) +def test_multi_output_getitem_is_numbered_by_escape_order(escaping): + graph = torch.fx.Graph() + x = graph.placeholder("x") + multi = graph.call_function( + torch.ops.aten.native_layer_norm.default, (x, [2], None, None, 1e-5) + ) + items = [graph.call_function(operator.getitem, (multi, i)) for i in range(3)] + graph.output(tuple(items[i] for i in escaping)) + visitor = _visitor(graph) + for i in range(3): + name = visitor.get_tensor_name(multi, i) + if i in escaping: + assert name == f"output_{escaping.index(i)}_{multi.name}@{i}" + else: + assert name == f"{multi.name}@{i}" + + +def test_direct_and_repeated_outputs_use_first_program_position(): + graph = torch.fx.Graph() + x = graph.placeholder("x") + y = graph.call_function(torch.neg, (x,)) + graph.output((y, y)) + assert _visitor(graph).get_tensor_name(y) == f"output_0_{y.name}@0" + + +def test_single_output_graph_is_numbered(): + graph = torch.fx.Graph() + x = graph.placeholder("x") + y = graph.call_function(torch.neg, (x,)) + graph.output(y) + assert _visitor(graph).get_tensor_name(y) == f"output_0_{y.name}@0" + + +def test_mutable_buffer_names_remain_unchanged(): + graph = torch.fx.Graph() + cache = graph.placeholder("b_cache") + x = graph.placeholder("x") + updated = graph.call_function(torch.add, (cache, x)) + ordinary = graph.call_function(torch.neg, (x,)) + graph.output((updated, ordinary)) + visitor = _visitor( + graph, buffers={cache.name: "cache"}, mutations={updated.name: "cache"} + ) + assert visitor.get_tensor_name(cache) == f"input_0_mutbuf_0_{cache.name}@0" + assert visitor.get_tensor_name(updated) == f"output_mutbuf_0_{updated.name}@0" + assert visitor.get_tensor_name(ordinary) == f"output_1_{ordinary.name}@0" + + +def test_op_validation_does_not_publish_output_names(): + graph = torch.fx.Graph() + x = graph.placeholder("x") + y = graph.call_function(torch.neg, (x,)) + graph.output((y,)) + assert _visitor(graph, validation=True).get_tensor_name(y) == f"{y.name}@0" diff --git a/backends/qualcomm/tests/test_partitioner_options.py b/backends/qualcomm/tests/test_partitioner_options.py new file mode 100644 index 00000000000..3286f9382e2 --- /dev/null +++ b/backends/qualcomm/tests/test_partitioner_options.py @@ -0,0 +1,73 @@ +# 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. + +import pytest +from executorch.backends.qualcomm.partition.qnn_partitioner import QnnPartitioner +from executorch.backends.qualcomm.serialization.qc_schema import ( + QnnExecuTorchBackendType, + QnnExecuTorchHtpPrecision, +) +from executorch.backends.qualcomm.serialization.qc_schema_serialize import ( + flatbuffer_to_option, +) +from executorch.backends.qualcomm.utils import utils + + +@pytest.mark.parametrize("backend", ["htp", "gpu", "lpai"]) +def test_from_options_matches_explicit_compiler_specs(backend): + if backend == "lpai" and utils.is_qnn_sdk_version_less_than("2.39"): + with pytest.raises(ValueError, match="requires QNN SDK version >= 2.39"): + QnnPartitioner.from_options("SM8850", backend=backend) + return + backend_options = {"use_fp16": True} if backend == "htp" else {} + partitioner = QnnPartitioner.from_options( + "SM8850", + backend=backend, + shared_buffer=True, + skip_node_id_set={"add"}, + skip_mutable_buffer=True, + **backend_options, + ) + specs = utils.generate_qnn_executorch_compiler_spec( + utils.get_soc_to_chipset_map()["SM8850"], + getattr(utils, f"generate_{backend}_compiler_spec")(**backend_options), + shared_buffer=True, + ) + assert [(s.key, s.value) for s in partitioner.compiler_specs] == [ + (s.key, s.value) for s in specs + ] + assert partitioner.skip_node_id_set == {"add"} + assert partitioner.skip_mutable_buffer + + +def test_from_options_routes_backend_and_partitioner_settings(): + partitioner = QnnPartitioner.from_options( + "SM8750", + use_fp16=True, + shared_buffer=True, + skip_node_op_set={"aten.add.Tensor"}, + convert_linear_to_conv2d=True, + ) + options = flatbuffer_to_option(partitioner.compiler_specs[0].value) + assert options.backend_options.backend_type == QnnExecuTorchBackendType.kHtpBackend + assert ( + options.backend_options.htp_options.precision + == QnnExecuTorchHtpPrecision.kHtpFp16 + ) + assert options.shared_buffer + assert partitioner.convert_linear_to_conv2d + assert partitioner.skip_node_op_set == {"aten.add.Tensor"} + + +@pytest.mark.parametrize("backend", ["htp", "gpu", "lpai"]) +def test_from_options_rejects_unknown_option(backend): + with pytest.raises(TypeError, match="Unknown QnnPartitioner option 'typo'"): + QnnPartitioner.from_options("SM8750", backend=backend, typo=True) + + +def test_from_options_rejects_unknown_backend(): + with pytest.raises(ValueError, match="Unsupported QNN backend"): + QnnPartitioner.from_options("SM8750", backend="invalid") diff --git a/backends/qualcomm/tests/test_passes.py b/backends/qualcomm/tests/test_passes.py index 606fed01d72..8ddbf5ca31c 100644 --- a/backends/qualcomm/tests/test_passes.py +++ b/backends/qualcomm/tests/test_passes.py @@ -28,7 +28,10 @@ ) from executorch.backends.qualcomm.builders.qnn_constants import OpContextLoader from executorch.backends.qualcomm.builders.utils import is_parameter -from executorch.backends.qualcomm.partition.qnn_partitioner import QnnOperatorSupport +from executorch.backends.qualcomm.partition.qnn_partitioner import ( + QnnOperatorSupport, + QnnPartitioner, +) from executorch.backends.qualcomm.qnn_preprocess import QnnBackend from executorch.backends.qualcomm.quantizer.quantizer import QnnQuantizer, QuantDtype from executorch.backends.qualcomm.serialization.qc_schema import ( @@ -850,6 +853,93 @@ def forward(self, x): self.skipTest(f"QNN SDK not available: {e}") raise + def test_partitioner_carries_qnn_edge_pass_pipeline(self): + """QnnPartitioner.transform_passes must reproduce the edge pass list that + to_edge_transform_and_lower_to_qnn previously built by hand, so a plain + to_edge_transform_and_lower(partitioner=[QnnPartitioner(...)]) lowers + through the same pipeline. + """ + + class Model(torch.nn.Module): + def __init__(self): + super().__init__() + self.linear = torch.nn.Linear(4, 4) + + def forward(self, x): + return torch.nn.functional.relu(self.linear(x)) + + sample_input = (torch.randn(1, 4),) + exported = torch.export.export(Model().eval(), sample_input, strict=True) + compiler_specs = generate_qnn_executorch_compiler_spec( + soc_model=QcomChipset.SM8650, + backend_options=generate_htp_compiler_spec(use_fp16=True), + ) + + partitioner = QnnPartitioner(compiler_specs) + expected = get_qnn_pass_manager_cls( + partitioner.backend + )().get_to_edge_transform_passes(exported, compiler_specs=compiler_specs) + actual = partitioner.transform_passes(exported) + + self.assertTrue(expected, "expected a non-empty QNN edge pass pipeline") + self.assertEqual([type(p) for p in actual], [type(p) for p in expected]) + + def test_partitioner_runs_qnn_aten_passes_pre_decomposition(self): + """transform_for_pre_decomposition must run the QNN ATen pipeline. + + DecomposeScaledDotProductAttention is part of that pipeline, so an SDPA + node must be gone after the hook runs. + """ + + class Sdpa(torch.nn.Module): + def forward(self, q, k, v): + return torch.nn.functional.scaled_dot_product_attention(q, k, v) + + sample_input = tuple(torch.randn(1, 2, 4, 8) for _ in range(3)) + exported = torch.export.export(Sdpa().eval(), sample_input, strict=True) + + def has_sdpa(program): + return any( + node.op == "call_function" + and "scaled_dot_product_attention" in str(node.target) + for node in program.graph.nodes + ) + + # Guard against a vacuous test. + self.assertTrue(has_sdpa(exported), "expected SDPA in the traced graph") + + compiler_specs = generate_qnn_executorch_compiler_spec( + soc_model=QcomChipset.SM8650, + backend_options=generate_htp_compiler_spec(use_fp16=True), + ) + transformed = QnnPartitioner(compiler_specs).transform_for_pre_decomposition( + exported + ) + self.assertFalse(has_sdpa(transformed)) + + def test_partitioner_from_options_matches_explicit_compiler_specs(self): + """from_options must build the same compiler specs a caller would + assemble by hand, and route partitioner-level options to __init__.""" + explicit = generate_qnn_executorch_compiler_spec( + soc_model=QcomChipset.SM8750, + backend_options=generate_htp_compiler_spec(use_fp16=True), + use_mha2sha=True, + ) + partitioner = QnnPartitioner.from_options( + soc_model="SM8750", + use_fp16=True, + use_mha2sha=True, + convert_linear_to_conv2d=True, + ) + self.assertEqual( + [(s.key, s.value) for s in partitioner.compiler_specs], + [(s.key, s.value) for s in explicit], + ) + self.assertTrue(partitioner.convert_linear_to_conv2d) + + with self.assertRaises(TypeError): + QnnPartitioner.from_options(soc_model="SM8750", not_an_option=1) + def test_index_put_int64_value_not_quantized(self): """QNN's IndexPut annotator must skip a non-float (int64) value arg. diff --git a/backends/qualcomm/utils/export_plan.py b/backends/qualcomm/utils/export_plan.py new file mode 100644 index 00000000000..7d242c6c126 --- /dev/null +++ b/backends/qualcomm/utils/export_plan.py @@ -0,0 +1,329 @@ +# 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. + +"""What a generic exporter needs from the QNN backend, in one object. + +An exporter (Hugging Face's, or any other) traces a model, optionally quantizes +it with a PT2E quantizer, and lowers it with ``to_edge_transform_and_lower``. +``QnnExportPlan`` fills in the QNN parts of that flow without the exporter +knowing about QNN: + +- ``program_transform``: after tracing and before quantization, a program with + a KV cache is rewritten into the KV-cache-as-I/O contract HTP needs + (``kv_cache_io``); anything else is left alone. +- ``calibration``: for a rewritten decoder, calibration runs the runner's token + loop over the calibration prompts; otherwise the exporter's own loop is used. +- ``partitioners``, ``method_name``, ``constant_methods``, ``backend_config``, + ``edge_compile_config``: the lowering settings, which depend on what the + transform produced. +""" + +from typing import Any, Callable, Dict, Iterable, List, Optional, Set + +import torch +from executorch.backends.qualcomm.partition.qnn_partitioner import QnnPartitioner +from executorch.backends.qualcomm.utils import kv_cache_io +from executorch.exir import EdgeCompileConfig +from executorch.exir.capture._config import ExecutorchBackendConfig +from executorch.exir.passes.memory_planning_pass import MemoryPlanningPass +from torch.export import ExportedProgram + + +class QnnExportPlan: + """QNN's part of a generic trace → quantize → lower export. + + Args: + soc_model: Target SoC, e.g. ``"SM8750"``. + backend: QNN backend, ``"htp"`` by default. + quantizer: The PT2E quantizer the program will be quantized with, or + None for an fp16 program. + metadata: Constant methods to carry for a decoder (``get_bos_id``, + ``get_eos_ids``), as a runtime reads them. + combine_methods: Lower a decoder's prompt and decode graphs into one + program (``prefill_forward`` + ``kv_forward``) that shares weights in + one QNN context, as ``qnn_llama_runner``'s hybrid mode loads it. + r3: For a decoder, rotate queries and keys by a Hadamard matrix after + RoPE (see ``to_kv_cache_io``). + seq_mse_candidates: For a quantized decoder, search each convolution's + weight scale over this many candidates after calibration (SeqMSE); + 0 skips it. + seq_mse_steps: A dict shared by the plans of graphs that share weights + (a decoder's prompt and decode steps): the first fills it, the + others reuse its weight scales. + masked_softmax: Mask attention with HTP's MaskedSoftmax form (see + ``to_kv_cache_io``). Defaults to on for quantized programs, where it + holds for any score range; an additive mask is faster on models + whose scores stay small. + **options: Further ``QnnPartitioner.from_options`` options. + """ + + def __init__( + self, + soc_model: str, + backend: str = "htp", + quantizer: Any = None, + metadata: Optional[Dict[str, Any]] = None, + combine_methods: bool = True, + r3: bool = False, + seq_mse_candidates: int = 0, + masked_softmax: Optional[bool] = None, + seq_mse_steps: Optional[Dict[str, float]] = None, + **options: Any, + ) -> None: + self.soc_model = soc_model + self.backend = backend + self.quantizer = quantizer + self.metadata = dict(metadata or {}) + self.options = dict(options) + if backend == "htp": + self.options.setdefault("use_fp16", quantizer is None) + self.kv_io: Optional[kv_cache_io.KvCacheIoProgram] = None + self.method_name = "forward" + self.combine_methods = combine_methods + self.r3 = r3 + self.seq_mse_candidates = seq_mse_candidates + self.seq_mse_steps = seq_mse_steps + self.masked_softmax = ( + quantizer is not None if masked_softmax is None else masked_softmax + ) + + # Trace → transform → quantize. + def program_transform(self, exported_program: ExportedProgram) -> ExportedProgram: + """Rewrite a program with a KV cache into the KV-cache-as-I/O contract.""" + if not kv_cache_io.has_kv_cache(exported_program): + return exported_program + self.kv_io = kv_cache_io.to_kv_cache_io( + exported_program, + self.metadata, + masked_softmax=self.masked_softmax, + r3=self.r3, + ) + self.method_name = self.kv_io.method_name + return self.kv_io.export() + + def calibration( + self, dataset: Optional[Iterable[Any]] + ) -> Optional[Callable[[torch.nn.Module], None]]: + """The token-loop calibration for a rewritten decoder, or None to let + the exporter calibrate as usual. ``dataset`` holds prompts as token-id + lists or as dicts with ``input_ids`` (and, for an encoder-decoder, + ``cross_kv``: the step's cross-attention inputs for that sample).""" + if self.kv_io is None or not dataset: + return None + prompts, cross_kv = [], [] + for sample in dataset: + ids = sample.get("input_ids") if isinstance(sample, dict) else sample + if isinstance(ids, torch.Tensor): + ids = ids.reshape(-1).tolist() + prompts.append(list(ids)) + if isinstance(sample, dict) and "cross_kv" in sample: + cross_kv.append(sample["cross_kv"]) + eos = self.metadata.get("get_eos_ids") or [] + return self.kv_io.calibration_fn( + prompts, + eos if isinstance(eos, list) else [eos], + seq_mse_candidates=self.seq_mse_candidates, + cross_kv=cross_kv or None, + seq_mse_steps=self.seq_mse_steps, + ) + + # Lowering. + def partitioners(self) -> List[QnnPartitioner]: + options = dict(self.options) + if self.kv_io is not None: + # HTP's decoder settings: multi-head attention split per head, and + # linear layers run as 1x1 convolutions. + options.setdefault("use_mha2sha", True) + options.setdefault("convert_linear_to_conv2d", True) + # The runner's KV buffers are shared (rpcmem) memory, and a prompt + # and a decode step compiled together share weights. + options.setdefault("shared_buffer", True) + options.setdefault("use_weight_sharing", True) + if self.quantizer is not None: + options.setdefault("passes_job", self._quant_io_passes()) + return [ + QnnPartitioner.from_options( + soc_model=self.soc_model, backend=self.backend, **options + ) + ] + + def group(self) -> Optional[str]: + """Graphs of one decoder are lowered together, so they share weights.""" + return "decoder" if self.kv_io is not None and self.combine_methods else None + + def merge(self, programs: Dict[str, ExportedProgram]) -> Dict[str, ExportedProgram]: + """Give every graph of the decoder the decode step's quantization + encodings: the runner hands the same KV buffers to all of them, and + weight sharing needs identical weight encodings.""" + if self.quantizer is None or kv_cache_io.METHOD_NAME not in programs: + return programs + source = programs[kv_cache_io.METHOD_NAME] + for name, program in programs.items(): + if name != kv_cache_io.METHOD_NAME: + share_encodings(source, program) + return programs + + def constant_methods( + self, program: Optional[ExportedProgram] = None + ) -> Optional[Dict[str, Any]]: + if self.kv_io is None: + return None + methods = self.kv_io.constant_methods() + if program is not None and self.quantizer is not None: + # The logits' quantization, so a runtime can dequantize them. + node = next(n for n in program.graph.nodes if n.op == "output").args[0][0] + while node.op == "call_function" and node.target not in _PER_TENSOR_QDQ: + node = node.args[0] + if node.op == "call_function": + methods["get_logits_scale"] = float(node.args[1]) + methods["get_logits_zero_point"] = int(node.args[2]) + return methods + + def edge_compile_config(self) -> EdgeCompileConfig: + # QNN's own passes emit non-core ATen ops. + return EdgeCompileConfig(_check_ir_validity=False) + + def backend_config( + self, alloc_graph_input: bool = True, alloc_graph_output: bool = True + ): + from executorch.backends.qualcomm._passes.build_quant_io import BuildQuantIo + + if self.kv_io is not None: + # The runtime binds the cache buffers itself. + alloc_graph_input = alloc_graph_output = False + return ExecutorchBackendConfig( + passes=[BuildQuantIo()], + memory_planning_pass=MemoryPlanningPass( + alloc_graph_input=alloc_graph_input, + alloc_graph_output=alloc_graph_output, + ), + ) + + def _quant_io_passes(self): + from executorch.backends.qualcomm._passes import TagQuantIO + from executorch.backends.qualcomm._passes.qnn_pass_manager import ( + get_qnn_pass_manager_cls, + ) + from executorch.backends.qualcomm.utils.constants import ( + QCOM_PASS_ACTIVATE_KEY, + QCOM_PASS_ARGS_KWARGS_DEFAULTS_KEY, + ) + + bits = _activation_bits(self.quantizer) + kv_dtype = torch.uint8 if bits == 8 else torch.uint16 + passes_job = get_qnn_pass_manager_cls().get_capture_program_passes() + passes_job[TagQuantIO][QCOM_PASS_ACTIVATE_KEY] = True + passes_job[TagQuantIO][QCOM_PASS_ARGS_KWARGS_DEFAULTS_KEY][ + "get_quant_io_dtype_fn" + ] = self.kv_io.quant_io_dtype_fn(kv_dtype, torch.uint16) + return passes_job + + +def _activation_bits(quantizer) -> int: + """Activation width of the quantizer's default config: the KV cache is an + activation, so this is the width it crosses the graph boundary at.""" + from executorch.backends.qualcomm.quantizer.custom_annotation import ( + annotate_kv_8bit, + ) + from executorch.backends.qualcomm.quantizer.quantizer import QuantDtype + + recipe = getattr(quantizer, "_recipe", None) + if recipe is not None: + if annotate_kv_8bit in ( + getattr(c, "func", c) for c in recipe.custom_quant_annotations + ): + return 8 + dtype = recipe._default_quant_dtype + else: + dtype = quantizer.default_quant_config.quant_dtype + return 8 if dtype in (QuantDtype.use_8a8w, QuantDtype.use_8a4w) else 16 + + +_PER_TENSOR_QDQ = ( + torch.ops.quantized_decomposed.quantize_per_tensor.default, + torch.ops.quantized_decomposed.dequantize_per_tensor.default, +) + + +def _encoding_key( + node: torch.fx.Node, output_index: Dict[torch.fx.Node, int], inputs: Set[str] +): + """Where a quantize/dequantize node's encoding comes from, comparable across + graphs of one model: a graph input by name, a graph output by position, + anything else by the op and module that produce the value. Constants + (numbered per graph) are left alone.""" + producer = node.args[0] + while producer.op == "call_function" and producer.target in _PER_TENSOR_QDQ: + producer = producer.args[0] + if producer.op == "placeholder" and producer.name in inputs: + return ("input", producer.name) + if producer.op == "placeholder": + # A constant's encoding follows from its value, the same in every graph. + return ("constant", producer.name) + users = [node] + while users: + user = users.pop() + if user in output_index: + return ("output", output_index[user]) + users.extend( + u + for u in user.users + if u.op == "call_function" and u.target in _PER_TENSOR_QDQ + ) + stack = producer.meta.get("nn_module_stack") or {} + path = list(stack.values())[-1][0] if stack else "" + return (str(producer.target), path) + + +def share_encodings( + source: ExportedProgram, target: ExportedProgram, io_only: bool = False +) -> int: + """Copy the per-tensor activation encodings of ``source`` onto ``target``, + two quantized graphs of the same model that differ in step length (a decode + step and a prompt step). Values are matched by where their encoding comes + from (graph input name, graph output position, or producing op and module, + in order of appearance); the graph inputs and outputs, which the runtime + exchanges between the two graphs, must all match. Returns the number of + encodings copied. With ``io_only``, only the graph inputs' and outputs' + (graphs whose insides differ, e.g. in their attention form).""" + + def keyed(program): + out = next(n for n in program.graph.nodes if n.op == "output") + output_index = { + arg: i + for i, arg in enumerate(out.args[0]) + if isinstance(arg, torch.fx.Node) + } + inputs = set(program.graph_signature.user_inputs) + groups: Dict[Any, List[torch.fx.Node]] = {} + for n in program.graph.nodes: + if n.op == "call_function" and n.target in _PER_TENSOR_QDQ: + groups.setdefault( + (_encoding_key(n, output_index, inputs), n.target), [] + ).append(n) + return groups + + src, dst = keyed(source), keyed(target) + copied = 0 + for key, nodes in dst.items(): + if io_only and key[0][0] not in ("input", "output"): + continue + peers = src.get(key) + if peers is not None and len({tuple(p.args[1:]) for p in peers}) == 1: + # One encoding however many nodes carry it (e.g. a dequantize per + # user of a graph input). + peers = peers[:1] * len(nodes) + if peers is None or len(peers) != len(nodes): + if key[0][0] in ("input", "output"): + raise RuntimeError( + f"Cannot share the encoding of graph I/O {key[0]}: no matching node." + ) + continue + for a, b in zip(peers, nodes): + b.args = (b.args[0], *a.args[1:]) + copied += 1 + target.graph_module.recompile() + return copied diff --git a/backends/qualcomm/utils/hf_export.py b/backends/qualcomm/utils/hf_export.py new file mode 100644 index 00000000000..099cb257bbb --- /dev/null +++ b/backends/qualcomm/utils/hf_export.py @@ -0,0 +1,351 @@ +# 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. + +"""Export a Hugging Face decoder LLM to a QNN ``.pte`` from ExecuTorch. + +Hugging Face exporters only trace: ``DynamoExporter.export_for_generation`` +with a static cache gives each generation step as a program that takes its +cache as inputs. Everything after that is ExecuTorch's and QNN's, through the +standard APIs: the KV-cache-as-I/O rewrite (``to_kv_cache_io``), PT2E +quantization with QNN's token-loop calibration, and +``to_edge_transform_and_lower``. By default the program is labelled +``kvcache_io``, which ExecuTorch's LLM runner (``KvCacheIOManager``) and +``qnn_llama_runner`` drive; with ``--io_contract hf_static`` it keeps Hugging +Face's static-cache ABI (``hf_static``), which ExecuTorch's LLM runner +drives through ``HfStaticIOManager``. + + python -m executorch.backends.qualcomm.utils.hf_export \\ + --model HuggingFaceTB/SmolLM2-135M-Instruct --soc_model SM8750 \\ + --quantize --hybrid --output smollm2.pte +""" + +import argparse +import copy +from typing import Any, Dict, Optional, Sequence + +import torch +from executorch.backends.qualcomm.utils.export_plan import QnnExportPlan +from executorch.exir import ExecutorchProgramManager, to_edge_transform_and_lower + +DEFAULT_CALIBRATION_PROMPTS = ( + "Once upon a time", + "The capital of France is", + "In a shocking finding, scientists discovered", + "The history of the Roman Empire began when", + "To make a good cup of coffee, you need", + "The stock market fell sharply on Monday after", + "Photosynthesis is the process by which plants", + "def fibonacci(n):", +) + + +def _quantize(program, quantizer, calibrate): + from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e + + args, kwargs = program.example_inputs + prepared = prepare_pt2e(program.module(), quantizer) + calibrate(prepared) + converted = convert_pt2e(prepared, fold_quantize=True) + return torch.export.export( + converted, copy.deepcopy(args), copy.deepcopy(kwargs), strict=False + ) + + +def export_hf_decoder( + model: torch.nn.Module, + tokenizer: Any, + soc_model: str, + *, + max_seq_len: int = 128, + prefill_len: Optional[int] = None, + quantizer: Any = None, + calibration_prompts: Sequence[str] = DEFAULT_CALIBRATION_PROMPTS, + io_contract: str = "kvcache_io", + **plan_options: Any, +) -> ExecutorchProgramManager: + """Export a decoder LLM to one QNN program. + + Args: + model: A Hugging Face causal LM. + tokenizer: Its tokenizer, for the sample and calibration prompts. + soc_model: Target SoC, e.g. ``"SM8750"``. + max_seq_len: Context length: the static cache's size. + prefill_len: With a length, the program also gets a prompt graph + taking that many tokens per step, sharing weights with the decode + graph. + quantizer: A PT2E QNN quantizer, or None for fp16. + calibration_prompts: Prompts QNN's calibration generates from. + io_contract: The program's cache contract. ``"kvcache_io"``: QNN's + KV cache as I/O (``prefill_forward``/``kv_forward``; the runner + writes the new rows). ``"hf_static"``: Hugging Face's + static-cache ABI unchanged (``prefill``/``decode``; the program writes + the new rows into the caller's caches, see ``hf_static``). + **plan_options: For ``kvcache_io``, ``QnnExportPlan`` options (``r3``, + ``seq_mse_candidates``, ``masked_softmax``, partitioner options); + for ``hf_static``, ``masked_softmax``, ``seq_mse_candidates`` + and partitioner options (not ``r3``: + it would change what the caches hold). + """ + if io_contract == "hf_static": + from executorch.backends.qualcomm.utils import hf_static + + if plan_options.pop("r3", False): + # R3 would change what the caches hold, which HF's ABI exposes. + raise ValueError("r3 is not supported with hf_static.") + return hf_static.export_hf_static( + model, + tokenizer, + soc_model, + max_seq_len=max_seq_len, + prefill_len=prefill_len, + quantizer=quantizer, + calibration_prompts=calibration_prompts, + **plan_options, + ) + if io_contract != "kvcache_io": + raise ValueError(f"Unknown io_contract {io_contract!r}.") + from transformers import GenerationConfig + from transformers.exporters import DynamoConfig, DynamoExporter + + if prefill_len: + words = tokenizer( + " ".join(["word"] * 2 * prefill_len), + return_tensors="pt", + add_special_tokens=False, + ) + sample = words["input_ids"][:, :prefill_len] + else: + # A short sample: the static cache's length must not grow to fit it. + sample = tokenizer(calibration_prompts[0], return_tensors="pt")["input_ids"] + sample = sample[:, : min(sample.shape[1], 8)] + generation = GenerationConfig( + do_sample=False, cache_implementation="static", max_cache_len=max_seq_len + ) + artifacts = DynamoExporter().export_for_generation( + model.eval(), + {"input_ids": sample, "attention_mask": torch.ones_like(sample)}, + DynamoConfig(), + generation_config=generation, + ) + + config = model.generation_config + eos = config.eos_token_id + metadata = { + "get_bos_id": config.bos_token_id, + "get_eos_ids": eos if isinstance(eos, list) or eos is None else [eos], + } + prompts = [tokenizer(p)["input_ids"] for p in calibration_prompts] + plans: Dict[str, QnnExportPlan] = {} + programs = {} + # The decode step searches weight scales once; the prompt step reuses them. + plan_options.setdefault("seq_mse_steps", {}) + for component in ("decode", "prefill") if prefill_len else ("decode",): + plan = QnnExportPlan( + soc_model, quantizer=quantizer, metadata=metadata, **plan_options + ) + program = plan.program_transform(artifacts.components[component].artifact) + if quantizer is not None: + program = _quantize(program, quantizer, plan.calibration(prompts)) + plans[plan.method_name] = plan + programs[plan.method_name] = program + + decode = plans["kv_forward"] + programs = decode.merge(programs) + constant_methods = {} + for method, plan in plans.items(): + constant_methods.update(plan.constant_methods(programs[method]) or {}) + edge = to_edge_transform_and_lower( + programs, + partitioner={method: plan.partitioners() for method, plan in plans.items()}, + compile_config=decode.edge_compile_config(), + constant_methods=constant_methods, + ) + return edge.to_executorch(config=decode.backend_config()) + + +class _CrossAttentionInputs(torch.nn.Module): + """An encoder-decoder's encoder, plus the cross-attention keys (transposed) + and values each decoder layer computes from its output: the fixed inputs + of the rewritten decode step, computed once per input.""" + + def __init__(self, model: torch.nn.Module, start_token: int): + super().__init__() + self.model = model + self.start_token = start_token + + def forward(self, features: torch.Tensor): + from transformers.cache_utils import DynamicCache, EncoderDecoderCache + + states = self.model.get_encoder()(features)[0] + cache = EncoderDecoderCache(DynamicCache(), DynamicCache()) + self.model.get_decoder()( + input_ids=torch.full((1, 1), self.start_token, dtype=torch.long), + encoder_hidden_states=states, + past_key_values=cache, + use_cache=True, + ) + layers = cache.cross_attention_cache.layers + return ( + *[layer.keys.transpose(2, 3).contiguous() for layer in layers], + *[layer.values for layer in layers], + ) + + +def export_hf_seq2seq( + model: torch.nn.Module, + soc_model: str, + features: Sequence[torch.Tensor], + *, + max_seq_len: int = 128, + quantizer: Any = None, + encoder_quantizer: Any = None, + **plan_options: Any, +) -> ExecutorchProgramManager: + """Export an encoder-decoder model (e.g. Whisper) to one QNN program. + + The program has an ``encoder`` method, returning the cross-attention keys + and values, and a ``kv_forward`` decode step that takes them as fixed + inputs after its self-attention caches (``kvcache_io`` with + ``get_n_cross_layers``). + + Args: + model: A Hugging Face encoder-decoder model. + soc_model: Target SoC, e.g. ``"SM8750"``. + features: Encoder inputs (e.g. log-mel features): the first is the + sample to trace with, all calibrate. + max_seq_len: Decoder context length: the static self-attention cache. + quantizer: A PT2E QNN quantizer for the decoder, or None for fp16. + encoder_quantizer: The encoder's quantizer (defaults to ``quantizer``). + **plan_options: ``QnnExportPlan`` options. + """ + from executorch.backends.qualcomm.partition.qnn_partitioner import QnnPartitioner + from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e + from transformers.exporters import DynamoConfig, DynamoExporter + + model = model.eval() + generation = copy.deepcopy(model.generation_config) + generation.do_sample = False + generation.cache_implementation = "static" + generation.max_cache_len = max_seq_len + generation.max_new_tokens = max_seq_len // 2 + artifacts = DynamoExporter().export_for_generation( + model, + {"input_features": features[0]}, + DynamoConfig(), + generation_config=generation, + ) + start = model.config.decoder_start_token_id + eos = generation.eos_token_id + metadata = { + "get_bos_id": start, + "get_eos_ids": eos if isinstance(eos, list) or eos is None else [eos], + } + plan = QnnExportPlan( + soc_model, quantizer=quantizer, metadata=metadata, **plan_options + ) + decode = plan.program_transform(artifacts.components["decode"].artifact) + + cross = _CrossAttentionInputs(model, start) + encoder = torch.export.export(cross, (features[0],), strict=False) + if quantizer is not None: + with torch.no_grad(): + samples = [cross(f) for f in features] + dataset = [{"input_ids": [start], "cross_kv": kv} for kv in samples] + decode = _quantize(decode, quantizer, plan.calibration(dataset)) + prepared = prepare_pt2e(encoder.module(), encoder_quantizer or quantizer) + for f in features: + prepared(f) + encoder = torch.export.export( + convert_pt2e(prepared, fold_quantize=True), (features[0],), strict=False + ) + edge = to_edge_transform_and_lower( + {"encoder": encoder, "kv_forward": decode}, + partitioner={ + "encoder": [ + QnnPartitioner.from_options( + soc_model=soc_model, + use_fp16=quantizer is None, + shared_buffer=True, + ) + ], + "kv_forward": plan.partitioners(), + }, + compile_config=plan.edge_compile_config(), + constant_methods=plan.constant_methods(decode), + ) + return edge.to_executorch(config=plan.backend_config()) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) + parser.add_argument("--model", required=True, help="Hugging Face model id or path") + parser.add_argument("--soc_model", default="SM8750") + parser.add_argument("--output", required=True, help="Path of the .pte") + parser.add_argument("--max_seq_len", type=int, default=128) + parser.add_argument( + "--hybrid", + action="store_true", + help="Add a prompt graph (prefill_forward) next to the decode graph", + ) + parser.add_argument("--prefill_len", type=int, default=32) + parser.add_argument( + "--quantize", + action="store_true", + help="Quantize with the model's tuned recipe (16-bit activations by default)", + ) + parser.add_argument( + "--recipe_id", default=None, help="Recipe to use instead of the model's" + ) + parser.add_argument( + "--io_contract", + choices=("kvcache_io", "hf_static"), + default="kvcache_io", + help="The program's cache contract (see export_hf_decoder)", + ) + args = parser.parse_args() + + from transformers import AutoModelForCausalLM, AutoTokenizer + + tokenizer = AutoTokenizer.from_pretrained(args.model) + model = AutoModelForCausalLM.from_pretrained(args.model, dtype=torch.float32) + quantizer, plan_options = None, {} + if args.quantize: + from executorch.backends.qualcomm.export_utils import ( + get_backend_type, + make_quantizer, + ) + from executorch.backends.qualcomm.quantizer.hf_llm_quant_recipe import ( + get_quant_recipe, + ) + + recipe = get_quant_recipe(args.recipe_id or args.model) + quantizer = make_quantizer( + backend=get_backend_type("htp"), soc_model=args.soc_model + ) + quantizer.set_recipe(recipe.recipe) + quantizer.set_convert_linear_to_conv2d(True) + plan_options = { + "masked_softmax": recipe.masked_softmax, + "seq_mse_candidates": recipe.seq_mse_candidates, + } + if args.io_contract == "kvcache_io": + plan_options["r3"] = recipe.r3 + program = export_hf_decoder( + model, + tokenizer, + args.soc_model, + max_seq_len=args.max_seq_len, + prefill_len=args.prefill_len if args.hybrid else None, + quantizer=quantizer, + io_contract=args.io_contract, + **plan_options, + ) + with open(args.output, "wb") as f: + program.write_to_file(f) + + +if __name__ == "__main__": + main() diff --git a/backends/qualcomm/utils/hf_static.py b/backends/qualcomm/utils/hf_static.py new file mode 100644 index 00000000000..50b753c6f9f --- /dev/null +++ b/backends/qualcomm/utils/hf_static.py @@ -0,0 +1,295 @@ +# 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. + +"""Lower the ``hf_static`` float-cache contract to QNN.""" + +from typing import Any, Optional, Sequence + +import torch +from executorch.backends.qualcomm.partition.qnn_partitioner import QnnPartitioner +from executorch.backends.qualcomm.utils import kv_cache_io +from executorch.backends.qualcomm.utils.export_plan import share_encodings +from executorch.backends.qualcomm.utils.hf_export import DEFAULT_CALIBRATION_PROMPTS +from executorch.exir import ( + EdgeCompileConfig, + ExecutorchBackendConfig, + to_edge_transform_and_lower, +) +from executorch.exir.passes import MemoryPlanningPass +from executorch.extension.llm.export.hf_static import ( + CPU_META, + DECODE_METHOD, + flat_call, + HfStaticLayout, + IO_CONTRACT, + IO_CONTRACT_VERSION, + PREFILL_METHOD, + sliding_windows, + specialize, + tag_cache_mutations, + to_hf_static, +) +from executorch.extension.llm.io_contracts.hf_static import HfStaticDriver +from torch.export import ExportedProgram + +try: + from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e +except ImportError: + convert_pt2e = prepare_pt2e = None + +try: + from transformers import GenerationConfig + from transformers.exporters import DynamoConfig, DynamoExporter +except ImportError: + DynamoExporter = None + + +def _cpu_nodes(program): + """Read explicit rewrite tags, including producers used only by CPU nodes.""" + if isinstance(program, ExportedProgram): + tag_cache_mutations(program) + cpu = {n for n in program.graph.nodes if n.meta.get("custom", {}).get(CPU_META)} + for node in reversed(list(program.graph.nodes)): + if node.users and all(user in cpu for user in node.users): + cpu.add(node) + return [n.name for n in program.graph.nodes if n in cpu] + + +class HfStaticPartitioner(QnnPartitioner): + """Keep tagged cache writes, masks and counter arithmetic outside QNN.""" + + def transform_passes(self, exported_program: ExportedProgram): + self.skip_node_id_set |= set(_cpu_nodes(exported_program)) + return super().transform_passes(exported_program) + + def partition(self, edge_program: ExportedProgram): + self.skip_node_id_set |= set(_cpu_nodes(edge_program)) + return super().partition(edge_program) + + +def _quantize( + programs, lengths, layout, quantizer, tokenizer, prompts, seq_mse_candidates=0 +): + """Quantize model compute with float caches and shared weight encodings.""" + if prepare_pt2e is None: + raise ImportError("hf_static quantization requires torchao PT2E.") + prepared = {m: prepare_pt2e(p.module(), quantizer) for m, p in programs.items()} + steps = {m: flat_call(programs[m], prepared[m]) for m in programs} + decode = HfStaticDriver({1: steps[DECODE_METHOD]}, layout) + for prompt in prompts: + ids = tokenizer(prompt)["input_ids"] + if len(ids) >= layout.max_seq_len: + raise ValueError( + f"Calibration prompt of {len(ids)} tokens doesn't fit the cache." + ) + sequence = decode.generate(ids, layout.max_seq_len - len(ids) - 1) + if PREFILL_METHOD in prepared: + chunk = lengths[PREFILL_METHOD] + prefill = HfStaticDriver( + {chunk: steps[PREFILL_METHOD], 1: steps[DECODE_METHOD]}, layout + ) + prefill.run(sequence) + if seq_mse_candidates: + first = tokenizer(prompts[0])["input_ids"] + scales = {} + + def run_decode(): + decode.generate(first, layout.max_seq_len - len(first) - 1) + + kv_cache_io.seq_mse( + prepared[DECODE_METHOD], seq_mse_candidates, run_decode, scales + ) + if PREFILL_METHOD in prepared: + kv_cache_io.seq_mse( + prepared[PREFILL_METHOD], seq_mse_candidates, lambda: None, scales + ) + quantized = {} + for method, program in programs.items(): + args, kwargs = program.example_inputs + with torch.fx.traceback.preserve_node_meta(): + quantized[method] = torch.export.export( + convert_pt2e(prepared[method], fold_quantize=True), + args, + kwargs, + strict=False, + ) + quantized[method].graph_module.meta.update(program.graph_module.meta) + for method in quantized: + if method != DECODE_METHOD: + share_encodings(quantized[DECODE_METHOD], quantized[method], io_only=True) + return quantized + + +def _attention_form(attention: str, query_len: int, quantized: bool) -> str: + """Auto uses split for decode or quantized compute, concat for float prefill.""" + if attention != "auto": + return attention + return "split" if query_len == 1 or quantized else "concat" + + +def _lower( + hf_graph, + lengths, + max_seq_len, + attention, + masked_softmax, + quantizer, + tokenizer, + calibration_prompts, + soc_model, + partitioner_options, + metadata, + seq_mse_candidates=0, + *, + config, +): + forms = { + m: _attention_form(attention, q, quantizer is not None) + for m, q in lengths.items() + } + programs = { + m: to_hf_static( + specialize(hf_graph, q, max_seq_len, config=config), + masked_softmax, + forms[m], + ) + for m, q in lengths.items() + } + layout = HfStaticLayout.of(programs[DECODE_METHOD]) + if layout.dtype not in (torch.float16, torch.float32): + raise ValueError("hf_static caches require float16 or float32.") + if quantizer is not None: + programs = _quantize( + programs, + lengths, + layout, + quantizer, + tokenizer, + calibration_prompts, + seq_mse_candidates, + ) + options = dict(partitioner_options) + options.setdefault("use_fp16", quantizer is None) + options.setdefault("shared_buffer", True) + options.setdefault("convert_linear_to_conv2d", True) + options.setdefault("use_weight_sharing", len(programs) > 1) + if "concat" in forms.values(): + options.setdefault("use_mha2sha", True) + partitioners = { + m: [HfStaticPartitioner.from_options(soc_model=soc_model, **options)] + for m in programs + } + constant_methods = { + **metadata, + "get_io_contract": IO_CONTRACT, + "get_io_contract_version": IO_CONTRACT_VERSION, + "get_max_seq_len": max_seq_len, + "get_n_layers": layout.n_layers, + "get_cache_dtype": str(layout.dtype).removeprefix("torch."), + "use_kv_cache": True, + } + edge = to_edge_transform_and_lower( + programs, + partitioner=partitioners, + compile_config=EdgeCompileConfig(_check_ir_validity=False), + constant_methods={k: v for k, v in constant_methods.items() if v is not None}, + ) + return edge, programs, layout + + +@torch.no_grad() +def export_hf_static( + model: torch.nn.Module, + tokenizer: Any, + soc_model: str, + *, + max_seq_len: int = 128, + prefill_len: Optional[int] = 32, + quantizer: Any = None, + calibration_prompts: Optional[Sequence[str]] = None, + masked_softmax: Optional[bool] = None, + attention: Optional[str] = None, + seq_mse_candidates: int = 0, + **partitioner_options: Any, +): + """Export QNN compute with caller-owned float caches and fixed prefill/decode. + + ``auto`` attention uses split for decode and quantized compute, concat for + float prefill. Prefill must fit every sliding window. The registered runner + pads prompt chunks and falls back to decode for a near-context tail. + """ + calibration_prompts = calibration_prompts or DEFAULT_CALIBRATION_PROMPTS + if ( + not isinstance(max_seq_len, int) + or isinstance(max_seq_len, bool) + or max_seq_len <= 0 + ): + raise ValueError("max_seq_len must be a positive integer.") + if prefill_len is not None and ( + not isinstance(prefill_len, int) + or isinstance(prefill_len, bool) + or not 1 < prefill_len <= max_seq_len + ): + raise ValueError( + "prefill_len must be None or an integer between 2 and max_seq_len." + ) + windows = sliding_windows(model.config) + minimum = min([max_seq_len, *[w for w in windows if w is not None]]) + if prefill_len is not None and prefill_len > minimum: + raise ValueError( + f"prefill_len ({prefill_len}) exceeds the smallest sliding cache length ({minimum})." + ) + attention = attention or "auto" + if attention not in ("split", "concat", "auto"): + raise ValueError(f"Unknown attention form {attention!r}.") + masked_softmax = quantizer is not None if masked_softmax is None else masked_softmax + if DynamoExporter is None: + raise ImportError("hf_static export requires transformers.exporters.") + sample = tokenizer(calibration_prompts[0], return_tensors="pt")["input_ids"][ + :, : min(8, minimum) + ] + artifacts = DynamoExporter().export_for_generation( + model.eval(), + {"input_ids": sample, "attention_mask": torch.ones_like(sample)}, + DynamoConfig(dynamic=True), + generation_config=GenerationConfig( + do_sample=False, cache_implementation="static", max_cache_len=max_seq_len + ), + keep_all_logits=True, + ) + lengths = {DECODE_METHOD: 1} + if prefill_len is not None: + lengths[PREFILL_METHOD] = prefill_len + config = model.generation_config + eos = config.eos_token_id + metadata = { + "get_vocab_size": model.config.vocab_size, + "get_bos_id": config.bos_token_id, + "get_eos_ids": eos if isinstance(eos, list) or eos is None else [eos], + } + edge, _, _ = _lower( + artifacts.components["decode"].artifact, + lengths, + max_seq_len, + attention, + masked_softmax, + quantizer, + tokenizer, + calibration_prompts, + soc_model, + partitioner_options, + metadata, + seq_mse_candidates, + config=model.config, + ) + return edge.to_executorch( + ExecutorchBackendConfig( + run_reinplace_pass=True, + memory_planning_pass=MemoryPlanningPass( + alloc_graph_input=False, alloc_graph_output=False + ), + ) + ) diff --git a/backends/qualcomm/utils/kv_cache_io.py b/backends/qualcomm/utils/kv_cache_io.py new file mode 100644 index 00000000000..bd04b26fc35 --- /dev/null +++ b/backends/qualcomm/utils/kv_cache_io.py @@ -0,0 +1,603 @@ +# 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. + +"""Rewrite an exported decoder step into QNN's KV-cache-as-I/O contract. + +Hugging Face exports a decoder step with its static cache as plain inputs: each +layer's full K and V cache goes in, the step writes the new token into it with +``index_copy_`` and attends over the whole buffer, and the updated buffers come +back out. HTP wants a different contract, the one ``qnn_llama_runner`` drives: + + (input_tokens, atten_mask, pos_ids, past_k_0..N-1, past_v_0..N-1) + -> (logits, new_k_0..N-1, new_v_0..N-1) + +where ``past_k_i`` is ``[B, H, D, S-1]`` (K stored transposed, so attention is a +plain matmul), ``past_v_i`` is ``[B, H, S-1, D]``, only the new token's K/V slice +is returned, and the mask is an additive float tensor managed by the runner. + +``to_kv_cache_io`` does that rewrite on the traced graph, after a normal export: +no model wrapper and nothing model-specific. It matches the cache writes and the +attention that reads them, so it applies to any decoder Hugging Face exports this +way, and it leaves models without a cache alone. +""" + +import functools +from dataclasses import dataclass, field +from typing import Any, Callable, Dict, List, Optional, Sequence + +import torch +from executorch.extension.llm.export.cache_graph import ( + fold_position_tables, + fold_weight_arithmetic, + norm_after_head_transpose, + producer_cache_write, + split_rotate_half, +) +from torch.export import ExportedProgram + +aten = torch.ops.aten + +# Label stored in the program so a runtime can pick the I/O manager that speaks +# this contract. +KV_CACHE_IO_CONTRACT = "kvcache_io" +KV_CACHE_IO_VERSION = 1 +METHOD_NAME = "kv_forward" +PREFILL_METHOD_NAME = "prefill_forward" +# Masked-slot value seen in calibration. Quantized, the runner writes code 0 +# for a masked slot and the mask's calibrated encoding decides what it means. +# At -65535 the softmax input's step is ~1.0: scores lose precision and HTP's +# softmax runs ~10x slower. -255 is what Qualcomm's native flow uses. +CALIBRATION_MASK_VALUE = -255.0 + +_CACHE_WRITES = (aten.index_copy_.default, aten.index_copy.default) +_SDPA = aten.scaled_dot_product_attention.default +# Steps SeqMSE searches over: about one calibration batch of Qualcomm's native +# flow, which searches on its first batch. +_SEQ_MSE_STEPS = 32 +# The R3 rotation's Hadamard matrix, shared by every layer. +_R3_WEIGHT = "_qnn_r3" +# Ops that only reshape or narrow a cache read between the write and the +# attention: the GQA head repeat, and the slice to the filled length some +# attention paths read (with ``is_causal``). The rewritten attention reads the +# whole window under the runner's mask instead, so both are dropped. +_REPEAT_OPS = ( + aten.slice.Tensor, + aten.unsqueeze.default, + aten.expand.default, + aten.reshape.default, + aten.view.default, + aten.clone.default, + aten._unsafe_view.default, +) + + +@dataclass +class KvCacheIoProgram: + """A decoder step rewritten into the QNN KV-cache-as-I/O contract.""" + + module: torch.fx.GraphModule + example_inputs: tuple + num_layers: int + max_seq_len: int + k_shapes: List[torch.Size] + v_shapes: List[torch.Size] + vocab_size: int + ar_len: int = 1 + logits_shape: Optional[torch.Size] = None + metadata: Dict[str, Any] = field(default_factory=dict) + # Encoder-decoder models: each cross-attention layer's keys (transposed, + # ``[1, H, D, S]``) and values (``[1, H, S, D]``), fixed for a whole + # generation, are inputs after the self-attention caches. + cross_k_shapes: List[torch.Size] = field(default_factory=list) + cross_v_shapes: List[torch.Size] = field(default_factory=list) + + @property + def method_name(self) -> str: + # The names qnn_llama_runner loads: a single-token decode step and the + # multi-token prompt step of hybrid mode. + return METHOD_NAME if self.ar_len == 1 else PREFILL_METHOD_NAME + + def export(self, strict: bool = False) -> ExportedProgram: + return torch.export.export(self.module, self.example_inputs, strict=strict) + + def constant_methods(self) -> Dict[str, Any]: + """Model metadata ``qnn_llama_runner`` reads, plus the contract label.""" + return { + **{k: v for k, v in self.metadata.items() if v is not None}, + "get_max_seq_len": self.max_seq_len, + "get_n_layers": self.num_layers, + **( + {"get_n_cross_layers": len(self.cross_k_shapes)} + if self.cross_k_shapes + else {} + ), + "get_vocab_size": self.vocab_size, + "use_kv_cache": True, + "get_io_contract": KV_CACHE_IO_CONTRACT, + "get_io_contract_version": KV_CACHE_IO_VERSION, + } + + def quant_io_dtype_fn(self, kv_dtype: torch.dtype, io_dtype: torch.dtype): + """``TagQuantIO`` callback: keep the cache, mask and logits quantized at the + graph boundary, which is the I/O the runner exchanges.""" + ar = self.ar_len + cache_in = {tuple(s) for s in self.k_shapes} | {tuple(s) for s in self.v_shapes} + cache_out = {(*s[:3], ar) for s in self.k_shapes} | { + (s[0], s[1], ar, s[3]) for s in self.v_shapes + } + mask_shape = tuple(self.example_inputs[1].shape) + logits_shape = tuple(self.logits_shape) + + def fn(node): + from executorch.backends.qualcomm.builders.utils import is_graph_output + + val = node.meta.get("val") + if not isinstance(val, torch.Tensor): + return None + shape = tuple(val.shape) + if node.op == "placeholder": + if shape in cache_in: + return kv_dtype + if shape == mask_shape and val.is_floating_point(): + return io_dtype + if is_graph_output(node): + if shape in cache_out: + return kv_dtype + if shape == logits_shape: + return io_dtype + return None + + return fn + + def calibration_fn( # noqa: C901 + self, + prompts: Sequence[Sequence[int]], + eos_ids: Sequence[int] = (), + seq_mse_candidates: int = 0, + cross_kv: Optional[Sequence[Sequence[torch.Tensor]]] = None, + seq_mse_steps: Optional[Dict[str, float]] = None, + ) -> Callable[[torch.nn.Module], None]: + """PT2E calibration that drives the step the way the runner does. + + Each prompt runs token by token; every step's new K/V slice is written + into the past buffers fed to the next step, and generation continues + greedily after the prompt until EOS or the context is full, so the + observers see real cache contents instead of an empty cache. + + ``cross_kv`` gives each prompt's cross-attention inputs (keys, then + values) for an encoder-decoder step. + + With ``seq_mse_candidates``, each convolution's weight scale is then + searched (SeqMSE) for the least output error over every step of the + first prompt, as Qualcomm's native flow does for its first batch. + ``seq_mse_steps`` shares the result between graphs (see ``seq_mse``). + """ + layers, ar, total = self.num_layers, self.ar_len, self.max_seq_len + past = total - ar + + def step(module, chunk, start, past_k, past_v, cross=()): + """Run ``chunk`` (at most ``ar`` tokens) at position ``start`` and + write its new K/V into the past buffers, as the runner does.""" + n = len(chunk) + mask = torch.full((1, 1, ar, total), CALIBRATION_MASK_VALUE) + for row in range(ar): + mask[..., row, :start] = 0.0 + mask[..., row, past : past + min(row, n - 1) + 1] = 0.0 + tokens = torch.tensor([chunk + [chunk[-1]] * (ar - n)], dtype=torch.int32) + positions = torch.tensor( + [[min(start + i, total - 1) for i in range(ar)]], dtype=torch.int32 + ) + outputs = module(tokens, mask, positions, *past_k, *past_v, *cross) + for i in range(layers): + past_k[i][..., start : start + n] = outputs[1 + i][..., :n] + past_v[i][..., start : start + n, :] = outputs[1 + layers + i][ + ..., :n, : + ] + return outputs[0] + + def run(module: torch.nn.Module, index: int) -> None: + tokens = list(prompts[index]) + cross = tuple(cross_kv[index]) if cross_kv else () + past_k = [torch.zeros(s) for s in self.k_shapes] + past_v = [torch.zeros(s) for s in self.v_shapes] + with torch.no_grad(): + pos = 0 + while pos < past: + if ar > 1: + # The prompt step only ever sees prompt chunks. + if pos >= len(tokens): + break + chunk = tokens[pos : min(pos + ar, past)] + step(module, chunk, pos, past_k, past_v, cross) + pos += len(chunk) + continue + logits = step(module, [tokens[pos]], pos, past_k, past_v, cross) + pos += 1 + if pos >= len(tokens): + nxt = int(torch.argmax(logits[0, -1]).item()) + if nxt in eos_ids: + break + tokens.append(nxt) + + def calibrate(module: torch.nn.Module) -> None: + for index in range(len(prompts)): + run(module, index) + if seq_mse_candidates and prompts: + seq_mse( + module, + seq_mse_candidates, + lambda: run(module, 0), + seq_mse_steps, + ) + + return calibrate + + +def _hadamard(n: int, dtype: torch.dtype) -> torch.Tensor: + """The ``n`` x ``n`` Hadamard matrix scaled to be orthonormal.""" + if n < 1 or n & (n - 1): + raise ValueError(f"R3 needs a power-of-two head dim, got {n}.") + h = torch.ones(1, 1, dtype=torch.float64) + while h.shape[0] < n: + h = torch.cat([torch.cat([h, h], 1), torch.cat([h, -h], 1)]) + return (h / n**0.5).to(dtype) + + +def seq_mse( + module: torch.fx.GraphModule, + num_candidates: int, + run: Callable[[], None], + steps: Optional[Dict[str, float]] = None, +) -> None: + """SeqMSE over every forward ``run`` makes: QNN's SeqMSE searches on each + forward it sees and keeps the last result, so with one token per step the + inputs are gathered first and searched once. + + ``steps`` maps each weight to its chosen scale step: filled when empty, + applied without searching otherwise, so graphs sharing weights (a prompt + and a decode step) quantize them identically.""" + from executorch.backends.qualcomm._passes.seq_mse import ( + InsertSeqMse, + RemoveSeqMse, + SeqMseModule, + ) + + # The rewritten program holds its weights as constants, not parameters. + module.get_parameter = lambda target: functools.reduce( + getattr, target.split("."), module + ) + try: + InsertSeqMse(num_candidates)(module) + finally: + del module.get_parameter + # Each search, keyed by its weight (conv -> weight observer -> weight). + searches = { + node.args[1].args[1].args[0].target: getattr(module, node.target) + for node in module.graph.nodes + if node.op == "call_module" + and isinstance(getattr(module, node.target), SeqMseModule) + } + # The Hadamard matrix's entries are all +-1/sqrt(d): any smaller scale + # would clip it. + searches.pop(_R3_WEIGHT, None) + if steps: + for weight, m in searches.items(): + m.best_candidate_step = steps.get(weight, 1.0) + else: + seen = {weight: ([], []) for weight in searches} + for weight, m in searches.items(): + m.forward = lambda x, y, io=seen[weight]: ( + len(io[0]) < _SEQ_MSE_STEPS + and (io[0].append(x.detach()), io[1].append(y.detach())) + ) + run() + for weight, m in searches.items(): + inputs, outputs = seen[weight] + if inputs: + m.best_candidate_step = m._find_best_candidate( + torch.cat(inputs), torch.cat(outputs) + ) + if steps is not None: + steps[weight] = m.best_candidate_step + RemoveSeqMse()(module) + + +def has_kv_cache(exported_program: ExportedProgram) -> bool: + """Whether the program writes a cache it takes as an input.""" + return any( + n.op == "call_function" and n.target in _CACHE_WRITES + for n in exported_program.graph.nodes + ) + + +def _source_input(node: torch.fx.Node) -> Optional[torch.fx.Node]: + """The program input a cache read comes from, if nothing writes it.""" + while node.op == "call_function" and node.target in _REPEAT_OPS: + node = node.args[0] + return node if node.op == "placeholder" else None + + +def _write_counters(layers) -> List[torch.fx.Node]: + """The scalar inputs (cache fill counters) the cache writes' indices derive + from.""" + counters, stack, seen = [], [], set() + for _, k_write, v_write in layers: + stack += [k_write.args[2], v_write.args[2]] + while stack: + node = stack.pop() + if not isinstance(node, torch.fx.Node) or node in seen: + continue + seen.add(node) + val = node.meta.get("val") + if node.op == "placeholder": + if isinstance(val, torch.Tensor) and val.dim() == 0: + counters.append(node) + continue + stack.extend(node.all_input_nodes) + return counters + + +def to_kv_cache_io( # noqa: C901 + exported_program: ExportedProgram, + metadata: Optional[Dict[str, Any]] = None, + masked_softmax: bool = False, + r3: bool = False, +) -> KvCacheIoProgram: + """Rewrite a Hugging Face decoder step into QNN's KV-cache-as-I/O contract. + + Args: + exported_program: A single-token decoder step exported with its static + cache as inputs (``past_key_values_*``), as Hugging Face exporters + produce it. + metadata: Constant-method metadata to carry, e.g. ``get_bos_id`` and + ``get_eos_ids``. + masked_softmax: Mask with ``where(mask == 0, scores, min(scores) - 20)`` + instead of adding the mask, which HTP fuses into its MaskedSoftmax + op. For quantized programs: unlike an additive mask, it holds + whatever the scores' range. + r3: Rotate queries and keys by a normalized Hadamard matrix after RoPE + (SpinQuant's R3, as Qualcomm's native flow applies it). Scores are + unchanged; outliers in the keys, and so in the cache, are spread + across the head dimension, which quantizes better. + + Returns: + A ``KvCacheIoProgram``: the rewritten module, its example inputs, and the + constant methods, quantized-I/O tagging and calibration loop that go with + the contract. + """ + gm = exported_program.module() + graph = gm.graph + placeholders = {n.name: n for n in graph.nodes if n.op == "placeholder"} + + # The attention calls, in layer order, and the cache writes they read. + # Attention reading caches that are only inputs (an encoder-decoder's + # cross-attention) takes those as fixed inputs. + layers, cross = [], [] + for node in graph.nodes: + if node.op == "call_function" and node.target is _SDPA: + k_write = producer_cache_write(node.args[1]) + v_write = producer_cache_write(node.args[2]) + if k_write is not None and v_write is not None: + layers.append((node, k_write, v_write)) + continue + k_in, v_in = _source_input(node.args[1]), _source_input(node.args[2]) + mask = node.args[3] if len(node.args) > 3 else node.kwargs.get("attn_mask") + if k_in is None or v_in is None or mask is not None: + raise ValueError( + f"{node.name} does not read a cache this rewrite understands." + ) + cross.append((node, k_in, v_in)) + if not layers: + raise ValueError("No attention reading a written cache was found.") + + query = layers[0][0].args[0].meta["val"] + if query.shape[0] != 1: + raise ValueError(f"Expected batch 1, got query {tuple(query.shape)}.") + ar = query.shape[2] + cache = layers[0][1].args[0].meta["val"] + max_seq_len = cache.shape[2] + dtype = cache.dtype + + tokens = placeholders.get("input_ids", placeholders.get("decoder_input_ids")) + position = placeholders.get("position_ids") + # Without position inputs, positions come from the caches' fill counters. + counters = [] if position is not None else _write_counters(layers) + if tokens is None or (position is None and not counters): + raise ValueError( + "Expected token ids and positions (position_ids or cache counters)." + ) + + first = next(iter(graph.nodes)) + with graph.inserting_before(first): + new_tokens = graph.placeholder("input_tokens") + new_mask = graph.placeholder("atten_mask") + new_position = graph.placeholder("pos_ids") + past_k, past_v = [], [] + for i in range(len(layers)): + past_k.append(graph.placeholder(f"past_k_{i}")) + for i in range(len(layers)): + past_v.append(graph.placeholder(f"past_v_{i}")) + cross_k = [graph.placeholder(f"cross_k_{i}") for i in range(len(cross))] + cross_v = [graph.placeholder(f"cross_v_{i}") for i in range(len(cross))] + + # Tokens and positions come in as int32, as the runner feeds them. + tokens.replace_all_uses_with(new_tokens) + if counters: + # Each counter holds the step's first position. + new_position.meta["val"] = torch.zeros(1, ar, dtype=torch.int32) + body = next(n for n in graph.nodes if n.op != "placeholder") + with graph.inserting_before(body): + row = graph.call_function(aten.select.int, (new_position, 0, 0)) + start = graph.call_function(aten.select.int, (row, 0, 0)) + for counter in counters: + value = graph.call_function( + aten._to_copy.default, + (start,), + {"dtype": counter.meta["val"].dtype}, + ) + counter.replace_all_uses_with(value) + position = new_position + fold_position_tables(gm, position, new_position, max_seq_len, dtype) + split_rotate_half(graph) + fold_weight_arithmetic(gm) + norm_after_head_transpose(graph) + if position is not new_position and position.users: + with graph.inserting_after(new_position): + widened = graph.call_function( + aten._to_copy.default, + (new_position,), + {"dtype": position.meta["val"].dtype}, + ) + position.replace_all_uses_with(widened) + + hadamard = None + if r3: + gm.register_buffer( + _R3_WEIGHT, _hadamard(cache.shape[3], dtype), persistent=False + ) + with graph.inserting_before(first): + hadamard = graph.get_attr(_R3_WEIGHT) + + k_shapes, v_shapes, new_ks, new_vs = [], [], [], [] + for i, (sdpa, k_write, v_write) in enumerate(layers): + q = sdpa.args[0] + k_new, v_new = k_write.args[3], v_write.args[3] + b, h_kv, s, d = k_write.args[0].meta["val"].shape + h_q = q.meta["val"].shape[1] + scale = sdpa.kwargs.get("scale", None) or d ** -0.5 + k_shapes.append(torch.Size([b, h_kv, d, s - ar])) + v_shapes.append(torch.Size([b, h_kv, s - ar, d])) + with graph.inserting_before(sdpa): + if hadamard is not None: + # As a linear layer, HTP runs it from the shared weights like the + # projections; a matmul with a constant operand runs much slower. + # H is symmetric, so x @ H == linear(x, H). + q = graph.call_function(aten.linear.default, (q, hadamard)) + k_new = graph.call_function(aten.linear.default, (k_new, hadamard)) + k_t = graph.call_function(aten.transpose.int, (k_new, 2, 3)) + k_all = graph.call_function(aten.cat.default, ([past_k[i], k_t], 3)) + v_all = graph.call_function(aten.cat.default, ([past_v[i], v_new], 2)) + if h_q != h_kv: + rep = h_q // h_kv + k_all = graph.call_function(aten.unsqueeze.default, (k_all, 2)) + k_all = graph.call_function( + aten.expand.default, (k_all, [b, h_kv, rep, d, s]) + ) + k_all = graph.call_function( + aten.reshape.default, (k_all, [b, h_q, d, s]) + ) + v_all = graph.call_function(aten.unsqueeze.default, (v_all, 2)) + v_all = graph.call_function( + aten.expand.default, (v_all, [b, h_kv, rep, s, d]) + ) + v_all = graph.call_function( + aten.reshape.default, (v_all, [b, h_q, s, d]) + ) + scores = graph.call_function(aten.matmul.default, (q, k_all)) + scores = graph.call_function(aten.mul.Tensor, (scores, scale)) + if masked_softmax: + low = graph.call_function(aten.amin.default, (scores, [-1], True)) + low = graph.call_function(aten.add.Tensor, (low, -20.0)) + visible = graph.call_function(aten.eq.Scalar, (new_mask, 0)) + scores = graph.call_function(aten.where.self, (visible, scores, low)) + else: + scores = graph.call_function(aten.add.Tensor, (scores, new_mask)) + probs = graph.call_function(aten._softmax.default, (scores, -1, False)) + attn = graph.call_function(aten.matmul.default, (probs, v_all)) + sdpa.replace_all_uses_with(attn) + new_ks.append(k_t) + new_vs.append(v_new) + + cross_k_shapes, cross_v_shapes = [], [] + for i, (sdpa, k_in, _) in enumerate(cross): + q = sdpa.args[0] + b, h_kv, s, d = k_in.meta["val"].shape + h_q = q.meta["val"].shape[1] + scale = sdpa.kwargs.get("scale", None) or d ** -0.5 + cross_k_shapes.append(torch.Size([b, h_kv, d, s])) + cross_v_shapes.append(torch.Size([b, h_kv, s, d])) + with graph.inserting_before(sdpa): + k_all, v_all = cross_k[i], cross_v[i] + if h_q != h_kv: + rep = h_q // h_kv + k_all = graph.call_function( + aten.repeat_interleave.self_int, (k_all, rep, 1) + ) + v_all = graph.call_function( + aten.repeat_interleave.self_int, (v_all, rep, 1) + ) + scores = graph.call_function(aten.matmul.default, (q, k_all)) + scores = graph.call_function(aten.mul.Tensor, (scores, scale)) + probs = graph.call_function(aten._softmax.default, (scores, -1, False)) + attn = graph.call_function(aten.matmul.default, (probs, v_all)) + sdpa.replace_all_uses_with(attn) + + # New outputs: logits, then each layer's new K slice, then each new V slice. + output = next(n for n in reversed(graph.nodes) if n.op == "output") + logits = ( + output.args[0][0] + if isinstance(output.args[0], (list, tuple)) + else output.args[0] + ) + logits_shape = list(logits.meta["val"].shape) + if ar > 1 and logits_shape[1] != ar: + # Hugging Face keeps only the last position's logits; the runner reads the + # row of the last real token in a padded prompt chunk, so keep them all. + node = logits + while node.op == "call_function" and node.args: + if node.target is aten.slice.Tensor and node.args[1] == 1: + node.replace_all_uses_with(node.args[0]) + logits_shape[1] = ar + break + node = node.args[0] + output.args = ((logits, *new_ks, *new_vs),) + + # Keep exactly what the new outputs are computed from. Everything else served + # the old contract only: the in-place cache writes, the cache-length counters + # and their updates, the old mask and the shape asserts on them. Some of those + # are side-effecting, so dead-code elimination alone would keep them. + live = set() + stack = [output] + while stack: + node = stack.pop() + if node in live: + continue + live.add(node) + stack.extend(node.all_input_nodes) + keep = {new_tokens, new_mask, new_position, *past_k, *past_v, *cross_k, *cross_v} + for node in reversed(list(graph.nodes)): + if node not in live and node not in keep: + graph.erase_node(node) + graph._codegen = torch.fx.graph.CodeGen() + graph.lint() + # A fresh module keeps only what the rewritten graph references, and none of + # the unlifted module's input-spec hooks for the old signature. + gm = torch.fx.GraphModule(gm, graph) + + vocab_size = logits.meta["val"].shape[-1] + example_inputs = ( + torch.zeros(1, ar, dtype=torch.int32), + torch.zeros(1, 1, ar, max_seq_len, dtype=dtype), + torch.zeros(1, ar, dtype=torch.int32), + *[torch.zeros(s, dtype=dtype) for s in k_shapes], + *[torch.zeros(s, dtype=dtype) for s in v_shapes], + *[torch.zeros(s, dtype=dtype) for s in cross_k_shapes], + *[torch.zeros(s, dtype=dtype) for s in cross_v_shapes], + ) + return KvCacheIoProgram( + module=gm, + example_inputs=example_inputs, + num_layers=len(layers), + max_seq_len=max_seq_len, + k_shapes=k_shapes, + v_shapes=v_shapes, + vocab_size=vocab_size, + ar_len=ar, + logits_shape=torch.Size(logits_shape), + metadata=dict(metadata or {}), + cross_k_shapes=cross_k_shapes, + cross_v_shapes=cross_v_shapes, + ) diff --git a/backends/qualcomm/utils/targets.bzl b/backends/qualcomm/utils/targets.bzl index 4ec50f0268b..98e4d459de9 100644 --- a/backends/qualcomm/utils/targets.bzl +++ b/backends/qualcomm/utils/targets.bzl @@ -16,5 +16,7 @@ def define_common_targets(): "//executorch/exir/backend:backend_details", "//executorch/exir/backend:compile_spec_schema", "//executorch/backends/qualcomm/serialization:serialization", + "//executorch/extension/llm/export:hf_static", + "//executorch/extension/llm/io_contracts:io_contracts", ], ) diff --git a/backends/qualcomm/utils/utils.py b/backends/qualcomm/utils/utils.py index 2f831a15b55..eea8831d246 100644 --- a/backends/qualcomm/utils/utils.py +++ b/backends/qualcomm/utils/utils.py @@ -462,9 +462,11 @@ def ensure_graph_specific_dict(value, graph_names): dep_table = ensure_graph_specific_dict(dep_table, graph_names) passes_job = ensure_graph_specific_dict(passes_job, graph_names) - # Prepare programs and partitioners + # Prepare programs and partitioners. The partitioner carries the QNN pass + # pipeline (transform_for_pre_decomposition / transform_passes), so lowering + # runs the same passes whether it is driven from here or from a plain + # to_edge_transform_and_lower call. aten_programs = {} - transform_passes = {} qnn_partitioners = { graph_name: [ QnnPartitioner( @@ -472,6 +474,9 @@ def ensure_graph_specific_dict(value, graph_names): skip_node_id_set=skip_node_id_set, skip_node_op_set=skip_node_op_set, skip_mutable_buffer=skip_mutable_buffer, + dep_table=dep_table[graph_name], + passes_job=passes_job[graph_name], + convert_linear_to_conv2d=convert_linear_to_conv2d, ) ] for graph_name in graph_names @@ -503,31 +508,10 @@ def ensure_graph_specific_dict(value, graph_names): dynamic_shapes=dynamic_shapes[graph_name], strict=True, ) - option = generate_qnn_executorch_option(compiler_specs[graph_name]) - python_options = flatbuffer_to_option(option) - backend_type = python_options.backend_options.backend_type - pass_manager = get_qnn_pass_manager_cls(backend_type)() - # This transformation is primarily intended for the LiftConstantScalarOperands pass - # to avoid creating temporary tensors in the operation builder. - # However, this pass will create a get_attr node, which should be converted - # into a lifted tensor constant by the lift_constant_tensor_pass. - # If placed in the to_edge_transform_passes, it will be executed - # after the lift_constant_tensor_pass, causing the operation builder - # to fail to correctly retrieve the parameter by the get_parameter. - aten_programs[graph_name] = pass_manager.transform_for_export_pipeline(ep) - transform_passes[graph_name] = pass_manager.get_to_edge_transform_passes( - ep, - passes_job=passes_job[graph_name], - dep_table=dep_table[graph_name], - compiler_specs=compiler_specs[graph_name], - skip_node_id_set=skip_node_id_set, - skip_node_op_set=skip_node_op_set, - convert_linear_to_conv2d=convert_linear_to_conv2d, - ) + aten_programs[graph_name] = ep with QnnManagerContext(compiler_specs): return to_edge_transform_and_lower( aten_programs, - transform_passes=transform_passes, partitioner=qnn_partitioners, constant_methods=constant_methods, compile_config=qnn_edge_config(), diff --git a/docs/source/backends-qualcomm.md b/docs/source/backends-qualcomm.md index c7d28c352d1..4f71e1eccb2 100644 --- a/docs/source/backends-qualcomm.md +++ b/docs/source/backends-qualcomm.md @@ -423,6 +423,36 @@ cd $env:EXECUTORCH_ROOT .\qnn_executor_runner.exe --model_path .\deeplab_v3\dlv3_qnn.pte ``` +### Export a Hugging Face LLM + +`executorch.backends.qualcomm.utils.hf_export` traces a Hugging Face decoder with HF's static-cache exporter and lowers it to QNN. The model can use a tuned quantization recipe while its cache I/O stays floating-point: + +```bash +python -m executorch.backends.qualcomm.utils.hf_export \ + --model Qwen/Qwen3-0.6B --soc_model SM8750 --max_seq_len 1024 \ + --quantize --hybrid --io_contract hf_static --output qwen3.pte +``` + +`--io_contract hf_static` produces `decode` (one token) and, with `--hybrid`, `prefill` (`--prefill_len` tokens), sharing weights and cache buffers. The program preserves HF's per-step tensor semantics and writes new cache rows on CPU; it does not preserve HF's single dynamic `forward` program interface. + +ExecuTorch's LLM runner uses `HfStaticIOManager`. The Python `hf_static` runner is registered by importing `executorch.extension.llm.io_contracts`; HF's `ExportedGenerator` can select it when transformers provides the I/O-contract hook. A runner must follow the [version-1 contract](llm-hf-static-io-contract.md), including fixed query lengths, padding rows in sliding masks, absolute counters and reset. This is not a claim that any stock HF-ABI runner can load the program. Python native bindings may clone returned outputs; a zero-copy Python host API is not claimed. + +#### `hf_static` and `kvcache_io` + +| | `hf_static` | `kvcache_io` | +|---|---|---| +| Step signature | HF static-cache tensor order; full K/V and counters are inputs and mutated outputs | Past K/V inputs; new K/V rows returned | +| Cache writes | Program CPU kernels into runner-owned buffers | Runner pastes returned rows | +| Methods | `prefill` / `decode` | `prefill_forward` / `kv_forward` | +| Runners | `HfStaticIOManager`; registered Python `hf_static` runner | `KvCacheIOManager`; `qnn_llama_runner`; registered Python adapter | +| Cache dtype | Float32 / float16 only, even with quantized model computation | Defined by the exported `kvcache_io` program | +| Auto attention | Decode uses split; prefill uses concat for floating-point model computation and split when quantized | QNN-specific cache layout and attention | +| Sliding layers | Window-sized caches with CPU shift / write; CPU correctness checked | Distinct cache-as-I/O contract; no equivalent sliding behavior is claimed | + +For the transformers export used here, the Python branch that shifts HF's sliding cache once full is traced away. `hf_static` replaces those writes and determines sliding layers from model configuration. A prompt method must not exceed the smallest sliding cache length. Padding rows in a window mask are all-false, so the shift counts real tokens only. See the [counter and padding protocol](llm-hf-static-io-contract.md#positions-counters-and-padding), including the counter difference from HF eager state. + +**Device boundary:** sliding-window shifting past context 512 has not been measured on device. CPU tests establish token correctness, not device support / latency. Quantized caches, `sha` prefill and full-length sliding-cache mode are held back from this series. + ### Run [Android LlamaDemo](https://github.com/meta-pytorch/executorch-examples/tree/main/llm/android/LlamaDemo) with QNN backend `$DEMO_APP` refers to the root of the executorch android demo, i.e., the directory containing `build.gradle.kts`. diff --git a/docs/source/index.md b/docs/source/index.md index 630a399e38d..bec73be31db 100644 --- a/docs/source/index.md +++ b/docs/source/index.md @@ -179,5 +179,6 @@ llm/working-with-llms advanced-topics-section tools-section api-section +llm-hf-static-io-contract support-section ``` diff --git a/docs/source/llm-hf-static-io-contract.md b/docs/source/llm-hf-static-io-contract.md new file mode 100644 index 00000000000..d83e7913c7c --- /dev/null +++ b/docs/source/llm-hf-static-io-contract.md @@ -0,0 +1,68 @@ +# Hugging Face static-cache I/O contract + +**Version:** 1 (`get_io_contract = "hf_static"`) + +## TL;DR + +An `hf_static` program keeps Hugging Face's per-step static-cache tensor signature. It is not a drop-in replacement for HF's single dynamic `forward` method: ExecuTorch exports fixed-length `prefill` / `decode` methods and requires the padding and counter protocol below. The runner owns the cache buffers; the program writes into them. + +## Goals / Non-Goals + +This contract defines the inputs, outputs, metadata and runner responsibilities for batch-one float-cache generation. It does not define runtime-owned dynamic caches, cache-state import / fork, quantized caches, or arbitrary query lengths. Sliding-window correctness is CPU-tested; this document does not claim sliding-window device performance past context 512. + +## Methods and metadata + +`decode` takes one token. An optional `prefill` method takes a fixed `N` tokens per call. Both methods share the same cache shapes, dtypes and buffers. If `prefill` is absent, prompts run through `decode`. + +Required metadata methods: + +| Method | Value / meaning | +|---|---| +| `get_io_contract` | `"hf_static"` | +| `get_io_contract_version` | `1`; a version-1 runner rejects newer versions | +| `get_max_seq_len` | Maximum total context length; not necessarily the longest physical cache when all layers slide | +| `get_n_layers` | Number of key / value / counter triples | +| `get_cache_dtype` | `"float32"` or `"float16"`, matching every cache input in every method | + +The exporter also records model vocabulary / BOS / EOS metadata when available and `use_kv_cache = True`. Runners validate contract metadata against method tensor metadata before allocating or writing buffers. Version 1 has no per-cache quantization scales or zero points. + +## Inputs and outputs + +Inputs are flattened in HF's exported order: + +1. `input_ids`: int64 `[1, N]` (`N = 1` for decode). +2. One or more bool attention masks: `[1, 1, N, W]`, in the exported order. +3. `position_ids`: int64 `[1, N]`. +4. Per layer, in layer order: key cache `[1, H_kv, L_i, D]`, value cache of the same shape and dtype, and a scalar int64 counter. + +Full-attention cache length `L_i` is the context length. A sliding-attention layer keeps `min(sliding_window, context_length)` rows. Sliding layers are determined from model configuration / HF cache layer types, not cache-size comparisons or mask names. Every layer's attention must be rewritten; an incomplete match is an export error. + +The exported Python module returns logits followed by the cache / counter state. The low-level ExecuTorch method prepends `3 * n_layers` user-input mutation outputs, then exposes logits at index `3 * n_layers`. The mutation prefix is opaque: its internal order need not match input order (for example, counters can precede K/V). Real HF graphs can also return the original `3 * n_layers` state tensors again after logits; runtimes accept either valid layout and do not treat that repeated state as a second owned cache. Mutation-prefix outputs write into their input buffers. The optional ordered state suffix may contain unplanned functional outputs (notably counters); the C++ runner binds each unplanned suffix output to the corresponding existing state buffer. Python native bindings may clone returned tensors, so callers must not infer program aliasing from returned Torch `data_ptr()` values. The C++ IOManager and registered Python host driver consume the fixed logits index and retain their existing state buffers. Python native bindings may clone returned outputs; those binding-owned copies are separate from the driver's persistent cache buffers and from the exported graph. The Python host API is not claimed to be zero-copy. In a padded chunk, Python may mark all padding rows false; the window-mask marking is mandatory, while unused full-attention padding rows can remain causal as in the C++ driver. Logits are floating-point `[1, N, vocabulary]`. For a padded prompt call, sample from row `real_tokens - 1`, not the last padded row. + +## Positions, counters and padding + +The runner's `start_pos` is the source of truth. Before every call, set every counter to `start_pos` and supply positions `start_pos .. start_pos + N - 1`. Real tokens must fit the total context. A fixed-length padded call must also fit physically; near the context boundary, run the remaining real tokens through `decode` instead of rejecting a prompt that fits. + +A prompt chunk with `r < N` real tokens repeats the last real token in its padding. Full-attention masks are causal by absolute position. In a sliding-window mask, each real query row must have at least one true entry and every padding row must be all-false. The program counts nonempty rows to determine `r`; this is a contract extension, not standard HF runner behavior. + +The program may advance counters by the static query length `N`, including padding. The runner must restore the logical count to `start_pos + r` after a padded call, or overwrite all counters with the next logical `start_pos` before execution. A reset returns the logical position and counters to zero. Unwritten cache slots need not contain zeros because masks exclude them. + +Sliding caches retain the most recent rows in chronological order. When a call crosses the window, the program shifts the retained rows and writes the real new rows on CPU. `prefill_len` must not exceed the smallest sliding cache length. Padding writes must use the last real row's data whenever they share its destination, so correctness does not depend on duplicate-index write order. + +HF's eager `StaticSlidingWindowLayer` may stop advancing its tensor `cumulative_length` once full, while its Python bookkeeping tracks the absolute position. This program's counters continue advancing as absolute positions. Cache state cannot be handed between HF eager and this contract by copying counters alone. + +## Runners and boundaries + +ExecuTorch's C++ `HfStaticIOManager` and the registered Python `hf_static` runner implement this protocol. The Python runner declares `drives_generation = True` for HF's `ExportedGenerator` contract hook. Importing `executorch.extension.llm.io_contracts` registers the adapters when the transformers checkout exposes that hook; a stock HF runner without it is not claimed to load these programs. Generation takes one unpadded sequence with standard absolute positions; caller-defined mask / position semantics are not supported. The runner's own fixed-chunk padding is supported. + +Cache allocation is backend-neutral. An optional I/O allocator can place buffers in QNN shared memory. Graph rewriting and CPU writes live in the LLM extension; QNN partitioning and HTP tuning remain in the Qualcomm backend. + +Cache inspection, adoption and forking are not exposed by the current IOManager. A future runtime-owned / forking contract needs its own state ownership API and version rather than silently changing version 1. + +## Version history + +| Version | Definition | +|---|---| +| 1 | Batch-one fixed-length decode / prefill, float caches, in-place state mutation, sliding-mask padding protocol | + +Experimental quantized caches, `sha` attention and full-length sliding-cache mode are excluded from this contract and the upstream series. diff --git a/examples/models/llama/CMakeLists.txt b/examples/models/llama/CMakeLists.txt index 443ccd01c7e..93d620587ad 100644 --- a/examples/models/llama/CMakeLists.txt +++ b/examples/models/llama/CMakeLists.txt @@ -175,6 +175,10 @@ if(TARGET qnn_executorch_backend) if(NOT MSVC) executorch_target_link_options_shared_lib(qnn_executorch_backend) endif() + if(TARGET qnn_llm_io_buffer_allocator) + list(APPEND link_libraries qnn_llm_io_buffer_allocator) + executorch_target_link_options_shared_lib(qnn_llm_io_buffer_allocator) + endif() endif() # MLX backend diff --git a/exir/backend/partitioner.py b/exir/backend/partitioner.py index c0edf0c71d8..286c39579a0 100644 --- a/exir/backend/partitioner.py +++ b/exir/backend/partitioner.py @@ -14,6 +14,7 @@ from executorch.exir.backend.backend_details import enforcedmethod from executorch.exir.backend.compile_spec_schema import CompileSpec +from executorch.exir.pass_manager import PassType from torch.export import ExportedProgram @@ -114,6 +115,31 @@ def ops_to_not_decompose( """ return ([], None) + def transform_passes(self, exported_program: ExportedProgram) -> List[PassType]: + """Return passes to run on the edge program before partitioning. + + EXIR invokes this hook automatically from + ``to_edge_transform_and_lower`` and runs the returned passes in the + same slot as its ``transform_passes`` argument, after any passes the + caller supplied there. Backend partitioners may override it to carry + edge-dialect rewrites that must run before ``partition`` sees the + graph, so callers do not have to thread backend-specific pass lists + through every lowering call. + + When multiple partitioners are used for a method, their passes are + concatenated in the order supplied to ``to_edge_transform_and_lower``. + + Args: + exported_program (ExportedProgram): The ATen-dialect program being + lowered, after ``transform_for_pre_decomposition`` has run and + before default decompositions. Supplied because pass + construction may depend on the graph. + + Returns: + List[PassType]: Passes to run, in order. Empty by default. + """ + return [] + def transform_for_pre_decomposition( self, exported_program: ExportedProgram ) -> ExportedProgram: diff --git a/exir/program/_program.py b/exir/program/_program.py index ea9b09b847f..2c46a92445f 100644 --- a/exir/program/_program.py +++ b/exir/program/_program.py @@ -245,6 +245,12 @@ def _create_empty_etrecord(): return ETRecord() +def _unused_name(prefix: str, index: int, taken: Set[str]) -> str: + while f"{prefix}{index}" in taken: + index += 1 + return f"{prefix}{index}" + + def lift_constant_tensor_pass(ep): """ Takes an ExportedProgram and returns the ExportedProgram modified in-place, @@ -255,6 +261,7 @@ def lift_constant_tensor_pass(ep): graph_signature = ep.graph_signature buffers = list(graph_signature.buffers) + taken = set(buffers) | set(ep.state_dict) | set(ep.constants) fake_mode = _detect_fake_mode_from_gm(ep.graph_module) @@ -271,7 +278,10 @@ def lift_constant_tensor_pass(ep): if not isinstance(constant_tensor, torch.Tensor): continue - constant_tensor_fqn = f"_lifted_tensor_constant{len(buffers)}" + constant_tensor_fqn = _unused_name( + "_lifted_tensor_constant", len(buffers), taken + ) + taken.add(constant_tensor_fqn) with ep.graph.inserting_before(first_user_input): # Insert the constant node before the first user input @@ -1131,12 +1141,13 @@ def _apply_pre_decomposition_transforms( return program -def _gen_edge_manager_for_partitioners( +def _gen_edge_manager_for_partitioners( # noqa: C901 partitioner: Dict[str, List[Partitioner]], aten_programs: Dict[str, ExportedProgram], config: EdgeCompileConfig, constant_methods: Optional[Dict[str, Any]], generate_etrecord: Optional[bool] = False, + collected_transform_passes: Optional[Dict[str, List[PassType]]] = None, ) -> "EdgeProgramManager": """ Generates EdgeProgramManager for subsequent lowering to the @@ -1157,6 +1168,16 @@ def _gen_edge_manager_for_partitioners( partitioners_for_program = partitioner.get(name, []) program = _apply_pre_decomposition_transforms(program, partitioners_for_program) + # Ask partitioners for edge passes here, while the program still matches + # what they saw in transform_for_pre_decomposition: pass construction can + # depend on the graph. + if collected_transform_passes is not None: + passes_for_program = [] + for curr_partitioner in partitioners_for_program: + passes_for_program.extend(curr_partitioner.transform_passes(program)) + if passes_for_program: + collected_transform_passes[name] = passes_for_program + # Functionalize program before asking partitioners to preserve ops program = program.run_decompositions({}) @@ -1285,6 +1306,46 @@ def collect_named_data_store_outputs( collect_named_data_store_outputs(exported_program.graph_module) +def _merge_partitioner_transform_passes( + transform_passes: Optional[ + Union[ + Sequence[PassType], + Dict[str, Sequence[PassType]], + GraphModulePassManager, + ExportedProgramPassManager, + ] + ], + partitioner_passes: Dict[str, List[PassType]], + method_names: Sequence[str], +) -> Dict[str, Sequence[PassType]]: + """Fold passes contributed by partitioners into ``transform_passes``. + + Caller-supplied passes run first, so an explicit argument is never + reordered behind a backend's own passes. + """ + if isinstance( + transform_passes, (GraphModulePassManager, ExportedProgramPassManager) + ): + raise ValueError( + "transform_passes was given as a PassManager, which cannot be " + "combined with passes contributed by a partitioner's " + "transform_passes(). Pass a list, or a dict of lists, instead." + ) + + merged: Dict[str, Sequence[PassType]] = {} + for name in method_names: + if isinstance(transform_passes, dict): + caller_passes = transform_passes.get(name, []) + elif transform_passes is not None: + caller_passes = transform_passes + else: + caller_passes = [] + combined = [*caller_passes, *partitioner_passes.get(name, [])] + if combined: + merged[name] = combined + return merged + + @et_logger("to_edge_transform_and_lower") def to_edge_transform_and_lower( # noqa: C901 programs: Union[ExportedProgram, Dict[str, ExportedProgram]], @@ -1363,10 +1424,21 @@ def to_edge_transform_and_lower( # noqa: C901 elif partitioner is None: partitioner = {name: [] for name in aten_programs.keys()} + partitioner_transform_passes: Dict[str, List[PassType]] = {} edge_manager = _gen_edge_manager_for_partitioners( - partitioner, aten_programs, config, constant_methods, generate_etrecord + partitioner, + aten_programs, + config, + constant_methods, + generate_etrecord, + collected_transform_passes=partitioner_transform_passes, ) + if partitioner_transform_passes: + transform_passes = _merge_partitioner_transform_passes( + transform_passes, partitioner_transform_passes, tuple(aten_programs) + ) + if transform_passes is not None: edge_manager = edge_manager.transform(transform_passes) for method in edge_manager.methods: diff --git a/exir/tests/targets.bzl b/exir/tests/targets.bzl index f68ee23cc16..7fc46ed309b 100644 --- a/exir/tests/targets.bzl +++ b/exir/tests/targets.bzl @@ -265,6 +265,20 @@ def define_common_targets(is_fbcode = False): ], ) + python_unittest( + name = "test_partitioner_transform_passes", + srcs = [ + "test_partitioner_transform_passes.py", + ], + deps = [ + "//caffe2:torch", + "//executorch/exir:lib", + "//executorch/exir:pass_base", + "//executorch/exir:pass_manager", + "//executorch/exir/backend:partitioner", + ], + ) + python_unittest( name = "common", srcs = [ diff --git a/exir/tests/test_partitioner_transform_passes.py b/exir/tests/test_partitioner_transform_passes.py new file mode 100644 index 00000000000..4fe404c22c0 --- /dev/null +++ b/exir/tests/test_partitioner_transform_passes.py @@ -0,0 +1,104 @@ +# 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. + +from typing import List + +import pytest + +import torch +from executorch.exir import to_edge_transform_and_lower +from executorch.exir.backend.partitioner import Partitioner, PartitionResult +from executorch.exir.pass_base import ExportPass, PassResult +from executorch.exir.pass_manager import PassType +from torch.export import export, ExportedProgram + + +class _Add(torch.nn.Module): + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x + x + + +class _RecordingPass(ExportPass): + def __init__(self, name: str, calls: List[str]) -> None: + super().__init__() + self.name = name + self.calls = calls + + def call(self, graph_module: torch.fx.GraphModule) -> PassResult: + self.calls.append(self.name) + return PassResult(graph_module, False) + + +class _ContributingPartitioner(Partitioner): + def __init__(self, names: List[str], calls: List[str]) -> None: + super().__init__() + self.passes = [_RecordingPass(name, calls) for name in names] + self.saw_program = False + + def transform_passes(self, exported_program: ExportedProgram) -> List[PassType]: + self.saw_program = isinstance(exported_program, ExportedProgram) + return list(self.passes) + + def partition(self, exported_program: ExportedProgram) -> PartitionResult: + return PartitionResult(exported_program, {}) + + +class _NoopPartitioner(Partitioner): + def partition(self, exported_program: ExportedProgram) -> PartitionResult: + return PartitionResult(exported_program, {}) + + +def _lower(partitioners, transform_passes=None): + exported_program = export(_Add(), (torch.randn(2, 3),), strict=True) + return to_edge_transform_and_lower( + exported_program, + transform_passes=transform_passes, + partitioner=partitioners, + ) + + +def test_partitioner_transform_passes_run() -> None: + calls: List[str] = [] + partitioner = _ContributingPartitioner(["a"], calls) + + _lower([partitioner]) + + assert calls == ["a"] + assert partitioner.saw_program + + +def test_partitioner_transform_passes_default_is_empty() -> None: + # A partitioner that does not override the hook must not change behavior. + _lower([_NoopPartitioner()]) + + +def test_partitioner_transform_passes_concatenate_in_order() -> None: + calls: List[str] = [] + first = _ContributingPartitioner(["first"], calls) + second = _ContributingPartitioner(["second"], calls) + + _lower([first, second]) + + assert calls == ["first", "second"] + + +def test_caller_transform_passes_run_before_partitioner_passes() -> None: + calls: List[str] = [] + partitioner = _ContributingPartitioner(["from_partitioner"], calls) + + _lower([partitioner], transform_passes=[_RecordingPass("from_caller", calls)]) + + assert calls == ["from_caller", "from_partitioner"] + + +def test_pass_manager_argument_with_contributed_passes_is_rejected() -> None: + from torch.fx.passes.infra.pass_manager import PassManager + + calls: List[str] = [] + partitioner = _ContributingPartitioner(["a"], calls) + + with pytest.raises(ValueError, match="cannot be combined"): + _lower([partitioner], transform_passes=PassManager([])) diff --git a/exir/tests/test_passes.py b/exir/tests/test_passes.py index 0ebd71c1690..ede501c6d97 100644 --- a/exir/tests/test_passes.py +++ b/exir/tests/test_passes.py @@ -1716,6 +1716,31 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: self.assertNotIn(name, new_ep.state_dict) self.assertIs(new_ep.state_dict["_prop_tensor_constant0"], stale) + def test_lift_constant_tensor_pass_avoids_buffer_name_collision(self) -> None: + class AddBuffer(torch.nn.Module): + def __init__(self): + super().__init__() + # A buffer named like a lifted constant, as re-exporting an + # unlifted program gives. + self.register_buffer("_lifted_tensor_constant1", torch.full((4,), 2.0)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x + self._lifted_tensor_constant1 + 3 + + edge = to_edge( + export(AddBuffer(), (torch.ones(4),), strict=True), + compile_config=EdgeCompileConfig(_skip_dim_order=False), + ) + edge = edge.transform([ScalarToTensorPass(), RemoveMixedTypeOperators()]) + exported_program = lift_constant_tensor_pass(edge.exported_program()) + + buffers = exported_program.graph_signature.buffers + self.assertEqual(len(buffers), len(set(buffers))) + torch.testing.assert_close( + exported_program.state_dict["_lifted_tensor_constant1"], + torch.full((4,), 2.0), + ) + def test_pass_no_user_inputs(self) -> None: class NoUserInputs(torch.nn.Module): def __init__(self): diff --git a/extension/llm/export/BUCK b/extension/llm/export/BUCK index d956bca2258..fa2f1e91f42 100644 --- a/extension/llm/export/BUCK +++ b/extension/llm/export/BUCK @@ -76,6 +76,25 @@ fbcode_target(_kind = runtime.python_library, ], ) +fbcode_target(_kind = runtime.python_library, + name = "hf_static", + srcs = [ + "cache_graph.py", + "hf_static.py", + ], + _is_external_target = True, + base_module = "executorch.extension.llm.export", + visibility = [ + "//executorch/backends/...", + "//executorch/examples/...", + "//executorch/extension/llm/...", + ], + deps = [ + "//caffe2:torch", + "//executorch/exir:lib", + ], +) + fbcode_target(_kind = runtime.python_library, name = "int4", srcs = [ diff --git a/extension/llm/export/cache_graph.py b/extension/llm/export/cache_graph.py new file mode 100644 index 00000000000..3e7c8ea4c69 --- /dev/null +++ b/extension/llm/export/cache_graph.py @@ -0,0 +1,414 @@ +# 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. + + +"""Shared cache-attention and position-table graph utilities.""" + +import functools +from typing import Any, Dict, List, Optional + +import torch + +aten = torch.ops.aten + +_CACHE_WRITES = (aten.index_copy_.default, aten.index_copy.default) + + +_WEIGHT_ARITHMETIC = ( + aten.add.Tensor, + aten.add.Scalar, + aten.mul.Tensor, + aten.mul.Scalar, + aten.rsub.Scalar, + aten._to_copy.default, + aten.to.dtype, + aten.to.dtype_layout, +) + + +_LAST_DIM_NORM = ( + aten.add.Tensor, + aten.add.Scalar, + aten.mul.Tensor, + aten.mul.Scalar, + aten.pow.Tensor_Scalar, + aten.rsqrt.default, + aten.mean.dim, + aten._to_copy.default, + aten.to.dtype, +) + + +_REPEAT_OPS = ( + aten.slice.Tensor, + aten.unsqueeze.default, + aten.expand.default, + aten.reshape.default, + aten.view.default, + aten.clone.default, + aten._unsafe_view.default, +) + + +def fold_weight_arithmetic(gm: torch.fx.GraphModule) -> None: + """Precompute values that depend only on weights, e.g. Gemma's RMSNorm + scale ``1 + weight``, so the normalization keeps the + ``x * rsqrt(mean(x^2) + eps) * weight`` shape QNN fuses into its RmsNorm + op. Unfused, the square overflows in fp16.""" + graph = gm.graph + count = 0 + for node in list(graph.nodes): + if node.op != "call_function" or node.target not in _WEIGHT_ARITHMETIC: + continue + tensors = [a for a in node.args if isinstance(a, torch.fx.Node)] + if not tensors or any(a.op != "get_attr" for a in tensors): + continue + args = [ + ( + functools.reduce(getattr, a.target.split("."), gm) + if isinstance(a, torch.fx.Node) + else a + ) + for a in node.args + ] + with torch.no_grad(): + value = node.target(*args, **node.kwargs) + name = f"_qnn_folded_{count}" + count += 1 + gm.register_buffer(name, value.contiguous(), persistent=False) + with graph.inserting_before(node): + folded = graph.get_attr(name) + folded.meta = dict(node.meta) + node.replace_all_uses_with(folded) + graph.erase_node(node) + + +def _on_last_dim(node: torch.fx.Node, dim: int) -> bool: + rank = node.meta["val"].dim() + return dim % rank == rank - 1 + + +def _rotate_half_input(cat: Any) -> Optional[tuple]: + """``(x, x1, x2)`` if ``cat`` is ``cat(-x2, x1)`` of ``x``'s halves.""" + if not isinstance(cat, torch.fx.Node) or cat.target != aten.cat.default: + return None + if len(cat.args[0]) != 2 or ( + len(cat.args) > 1 and not _on_last_dim(cat, cat.args[1]) + ): + return None + neg, x1 = cat.args[0] + if neg.target != aten.neg.default: + return None + x2 = neg.args[0] + if x1.target != aten.slice.Tensor or x2.target != aten.slice.Tensor: + return None + x = x1.args[0] + if x2.args[0] is not x or not _on_last_dim(x, x1.args[1]): + return None + half = x.meta["val"].shape[-1] // 2 + if tuple(x1.args[2:4]) != (0, half) or x2.args[2] != half: + return None + return x, x1, x2 + + +def _rope_in_halves(graph, x1, x2, cos, sin, half, tables) -> torch.fx.Node: + def part(t, lo, hi): + # Every layer reads the same cos/sin halves: slice each table once. + if (t, lo, hi) not in tables: + with graph.inserting_after(t): + tables[(t, lo, hi)] = graph.call_function( + aten.slice.Tensor, (t, -1, lo, hi) + ) + return tables[(t, lo, hi)] + + def mul(a, b): + return graph.call_function(aten.mul.Tensor, (a, b)) + + cos1, cos2 = part(cos, 0, half), part(cos, half, 2 * half) + sin1, sin2 = part(sin, 0, half), part(sin, half, 2 * half) + first = graph.call_function(aten.sub.Tensor, (mul(x1, cos1), mul(x2, sin1))) + second = graph.call_function(aten.add.Tensor, (mul(x2, cos2), mul(x1, sin2))) + return graph.call_function(aten.cat.default, ([first, second], -1)) + + +def split_rotate_half(graph: torch.fx.Graph) -> None: + """Rewrite RoPE written with ``rotate_half``, ``x * cos + cat(-x2, x1) * sin``, + into its two halves, ``cat(x1 * cos1 - x2 * sin1, x2 * cos2 + x1 * sin2)``. + + Same values; HTP runs the halved form faster, as it skips building the + rotated copy of ``x`` (the form Qualcomm's static_llama uses).""" + tables: Dict[tuple, torch.fx.Node] = {} + for add in list(graph.nodes): + if add.target != aten.add.Tensor or not all( + isinstance(a, torch.fx.Node) and a.target == aten.mul.Tensor + for a in add.args[:2] + ): + continue + for direct, rotated in (add.args[:2], add.args[1::-1]): + x, cos = direct.args + cat, sin = rotated.args + found = _rotate_half_input(cat) + if found is None or found[0] is not x: + continue + with graph.inserting_before(add): + new = _rope_in_halves( + graph, + found[1], + found[2], + cos, + sin, + x.meta["val"].shape[-1] // 2, + tables, + ) + new.meta["val"] = add.meta.get("val") + add.replace_all_uses_with(new) + break + + +def _rank(node: torch.fx.Node) -> int: + val = node.meta.get("val") + return val.dim() if isinstance(val, torch.Tensor) else 0 + + +def _norm_chain(out: torch.fx.Node) -> Optional[tuple]: + """``(source, nodes)`` if ``out`` normalizes one 4-D ``source`` over its + last dimension, through ops of ``_LAST_DIM_NORM`` and weights of at most + one dimension, and nothing else reads the intermediate values.""" + if out.op != "call_function" or out.target not in _LAST_DIM_NORM: + return None + nodes, stack, source = [], [out], None + while stack: + node = stack.pop() + if node in nodes: + continue + if node.op == "call_function" and node.target in _LAST_DIM_NORM: + if node.target == aten.mean.dim and not all( + _on_last_dim(node.args[0], d) for d in node.args[1] + ): + return None + nodes.append(node) + stack.extend(a for a in node.args if isinstance(a, torch.fx.Node)) + elif _rank(node) == 4 and source in (None, node): + source = node + elif _rank(node) > 1: + return None + if source is None or not any(n.target == aten.mean.dim for n in nodes): + return None + if not _only_read_within(nodes, out): + return None + return source, nodes + + +def _only_read_within(nodes: List[torch.fx.Node], out: torch.fx.Node) -> bool: + """Whether every value of ``nodes`` but ``out`` is read only by ``nodes`` + (shape assertions aside).""" + return all( + user in nodes or user.target == aten._assert_tensor_metadata.default + for node in nodes + if node is not out + for user in node.users + ) + + +def norm_after_head_transpose(graph: torch.fx.Graph) -> None: + """Move a per-head normalization that runs before the ``[B, S, H, D]`` → + ``[B, H, S, D]`` transpose to after it, e.g. Qwen3's query and key norms. + + Same values, as the norm is over ``D``; HTP runs it per head after the + attention is split by head, which the ``[B, S, H, D]`` layout turns into + strided copies (the order Qualcomm's static_llama uses).""" + for transpose in list(graph.nodes): + if transpose.target != aten.transpose.int: + continue + out = transpose.args[0] + dims = sorted(d % 4 for d in transpose.args[1:3]) + if out.meta["val"].dim() != 4 or dims != [1, 2] or len(out.users) != 1: + continue + found = _norm_chain(out) + if found is None: + continue + source, nodes = found + with graph.inserting_after(source): + moved = graph.call_function( + aten.transpose.int, (source, *transpose.args[1:3]) + ) + for node in nodes: + for user in list(node.users): + if user.target == aten._assert_tensor_metadata.default: + graph.erase_node(user) + node.replace_input_with(source, moved) + transpose.replace_all_uses_with(out) + graph.erase_node(transpose) + + +def producer_cache_write(node: torch.fx.Node) -> Optional[torch.fx.Node]: + """Walk back from an attention K/V argument to the cache write it reads.""" + while node.op == "call_function": + if node.target in _CACHE_WRITES: + return node + if node.target not in _REPEAT_OPS: + return None + node = node.args[0] + return None + + +def _position_only_nodes(gm: torch.fx.GraphModule, position: torch.fx.Node): + """Nodes computed from the position input and constants alone.""" + position_only = set() + for node in gm.graph.nodes: + if node is position or node.op == "get_attr": + position_only.add(node) + elif node.op in ("call_function", "call_module"): + inputs = node.all_input_nodes + if ( + inputs + and all(i in position_only for i in inputs) + and any( + i is position or i in position_only and i.op != "get_attr" + for i in inputs + ) + ): + position_only.add(node) + return position_only + + +def _eval_nodes(gm, nodes, position, value): + env = {} + for node in gm.graph.nodes: + if node not in nodes: + continue + if node is position: + env[node] = value + elif node.op == "get_attr": + obj = gm + for atom in node.target.split("."): + obj = getattr(obj, atom) + env[node] = obj + else: + args = torch.fx.node.map_arg(node.args, env.__getitem__) + kwargs = torch.fx.node.map_arg(node.kwargs, env.__getitem__) + if node.op == "call_module": + env[node] = gm.get_submodule(node.target)(*args, **kwargs) + else: + env[node] = node.target(*args, **kwargs) + return env + + +def fold_position_tables( # noqa: C901 + gm, position, new_position, max_seq_len, dtype, index_meta=None +): + """Replace every value that depends only on the position (RoPE cos/sin) with a + lookup into a table precomputed for all ``max_seq_len`` positions. + + HTP has no efficient path for the trigonometry, and a 16-bit quantized angle + loses precision quickly, so this is what QNN's own decoders do by hand. + """ + position_only = _position_only_nodes(gm, position) + frontier = [ + n + for n in position_only + if n.op != "get_attr" + and n is not position + and any(u not in position_only for u in n.users) + and isinstance(n.meta.get("val"), torch.Tensor) + ] + if not frontier: + return + # Fold at the source of each value rather than at every per-layer view of it, + # so all layers share one lookup. + roots = [] + for node in frontier: + while ( + node.target in _REPEAT_OPS + and node.args[0] in position_only + and node.args[0] is not position + and node.args[0].op != "get_attr" + and isinstance(node.args[0].meta.get("val"), torch.Tensor) + ): + node = node.args[0] + if node not in roots: + roots.append(node) + frontier = roots + pos_val = position.meta["val"] + ar = pos_val.shape[-1] + # Every position of the step at once: values that depend on the position + # element-wise (RoPE) are the same along the step axis then. + rows = [ + _eval_nodes( + gm, + position_only, + position, + torch.full(tuple(pos_val.shape), p, dtype=pos_val.dtype), + ) + for p in range(max_seq_len) + ] + graph = gm.graph + # Build and check every table on the untouched graph first, then rewrite. + tables: List[torch.Tensor] = [] + plans = [] + probe = torch.arange(ar, dtype=pos_val.dtype).reshape(pos_val.shape) + if ar < max_seq_len: + probe = probe + 1 + probe_env = _eval_nodes(gm, position_only, position, probe) if ar > 1 else None + for node in frontier: + out_shape = list(node.meta["val"].shape) + if ar == 1: + step_axis = None + entries = [row[node] for row in rows] + else: + if ar not in out_shape: + raise ValueError( + f"Cannot fold {node.name}: no axis matches the {ar} step positions." + ) + step_axis = out_shape.index(ar) + entries = [row[node].select(step_axis, 0) for row in rows] + table = torch.stack(entries).to(dtype).contiguous() + if step_axis is not None: + # The lookup must reproduce the value at mixed positions. + got = ( + table.index_select(0, probe.reshape(-1).long()) + .movedim(0, step_axis) + .reshape(out_shape) + ) + if not torch.allclose(got, probe_env[node].to(dtype), atol=1e-5): + raise ValueError( + f"{node.name} is not element-wise in the position; cannot fold it." + ) + index = next( + ( + j + for j, t in enumerate(tables) + if t.shape == table.shape and torch.equal(t, table) + ), + None, + ) + if index is None: + index = len(tables) + tables.append(table) + gm.register_buffer(f"_qnn_position_table_{index}", table, persistent=False) + plans.append((node, index, step_axis, out_shape, table.dim())) + for node, index, step_axis, out_shape, table_dim in plans: + with graph.inserting_before(node): + attr = graph.get_attr(f"_qnn_position_table_{index}") + flat = graph.call_function(aten.view.default, (new_position, [-1])) + if index_meta is not None: + flat.meta.update(index_meta) + picked = graph.call_function(aten.index_select.default, (attr, 0, flat)) + if step_axis is not None and step_axis != 0: + perm = list(range(1, table_dim)) + perm.insert(step_axis, 0) + picked = graph.call_function(aten.permute.default, (picked, perm)) + lookup = graph.call_function(aten.reshape.default, (picked, out_shape)) + node.replace_all_uses_with(lookup) + # Shape asserts on the folded values would keep the old subgraph alive. + for node in list(graph.nodes): + if ( + node.op == "call_function" + and node.target + in (aten._assert_tensor_metadata.default, aten._assert_scalar.default) + and all(i in position_only for i in node.all_input_nodes) + ): + graph.erase_node(node) diff --git a/extension/llm/export/hf_static.py b/extension/llm/export/hf_static.py new file mode 100644 index 00000000000..6526859b759 --- /dev/null +++ b/extension/llm/export/hf_static.py @@ -0,0 +1,708 @@ +# 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. + + +"""Backend-neutral rewrites of Hugging Face's per-step static-cache signature. + +The program has fixed ``prefill`` and ``decode`` methods. Runners supply masks, +absolute positions and cache counters; padding rows in window masks are false. +The program attends over incoming caches and new rows, then updates the caches. +""" + +import copy +from dataclasses import dataclass, field +from typing import Dict, List, Optional, Sequence + +import torch +import torch.utils._pytree as pytree +from executorch.exir.operator.convert import unwrap_op_overload +from executorch.extension.llm.export import cache_graph +from torch.export import ExportedProgram + +aten = torch.ops.aten +IO_CONTRACT = "hf_static" +# 1: float caches, fixed prefill/decode methods, absolute counters and padding masks. +IO_CONTRACT_VERSION = 1 +PREFILL_METHOD = "prefill" +DECODE_METHOD = "decode" +CACHE_PREFIX = "past_key_values_" +CPU_META = "hf_static_cpu" +_MASK_FILL = -10000.0 + + +@dataclass +class HfStaticLayout: + """What a runner and the calibration loop need to know about the inputs.""" + + input_names: List[str] + query_len: int + max_seq_len: int + n_layers: int + cache_shapes: List[torch.Size] + # The cache inputs' shapes in order: per layer keys, values and counter (rank 0). + state_shapes: List[torch.Size] + dtype: torch.dtype + # Each attention mask's width: the length of the caches it masks. + mask_widths: Dict[str, int] + state_dtypes: List[torch.dtype] = field(default_factory=list) + token_dtype: torch.dtype = torch.int64 + position_dtype: torch.dtype = torch.int64 + window_masks: List[str] = field(default_factory=list) + + @classmethod + def of(cls, exported_program: ExportedProgram) -> "HfStaticLayout": + names, vals = [], [] + for spec in exported_program.graph_signature.input_specs: + if spec.kind.name != "USER_INPUT": + continue + node = exported_program.graph.find_nodes( + op="placeholder", target=spec.arg.name + )[0] + names.append(spec.arg.name) + vals.append(node.meta["val"]) + state = [v for n, v in zip(names, vals) if n.startswith(CACHE_PREFIX)] + caches = [v for v in state if v.dim() == 4] + if len(state) != 3 * (len(caches) // 2) or any( + v.dim() not in (0, 4) for v in state + ): + raise ValueError( + "Expected per layer keys, values and a counter as the cache inputs." + ) + context = exported_program.graph_module.meta.get( + "hf_static_max_seq_len", max(c.shape[2] for c in caches) + ) + return cls( + input_names=names, + query_len=vals[names.index("input_ids")].shape[1], + max_seq_len=context, + n_layers=len(caches) // 2, + cache_shapes=[c.shape for c in caches], + state_shapes=[v.shape for v in state], + dtype=caches[0].dtype, + mask_widths={ + n: v.shape[-1] + for n, v in zip(names, vals) + if n.startswith("attention_mask") + }, + state_dtypes=[v.dtype for v in state], + token_dtype=vals[names.index("input_ids")].dtype, + position_dtype=vals[names.index("position_ids")].dtype, + window_masks=exported_program.graph_module.meta.get( + "hf_static_window_masks", [] + ), + ) + + +def sliding_windows(config) -> List[Optional[int]]: + """Per-layer windows from the model config, never inferred from cache size.""" + layers = getattr(config, "per_layer_config", [config] * config.num_hidden_layers) + types = getattr(config, "layer_types", None) + if types is not None and len(types) != len(layers): + raise ValueError("layer_types must describe every cache layer.") + windows = [] + for i, layer in enumerate(layers): + kind = types[i] if types is not None else None + window = getattr(layer, "sliding_window", None) + if kind == "full_attention": + window = None + elif kind not in (None, "sliding_attention"): + raise ValueError(f"Unsupported cache layer type {kind!r}.") + elif kind == "sliding_attention" and window is None: + window = getattr(config, "sliding_window", None) + if window is None: + raise ValueError("A sliding layer requires sliding_window.") + if window is not None and window <= 0: + raise ValueError("sliding_window must be positive.") + windows.append(window) + return windows + + +def specialize( # noqa: C901 + exported_program: ExportedProgram, + query_len: int, + max_seq_len: int, + *, + config, +) -> ExportedProgram: + """Re-export at fixed lengths, preserving config-defined sliding caches.""" + if not 0 < query_len <= max_seq_len: + raise ValueError("query_len must be positive and no longer than max_seq_len.") + windows = sliding_windows(config) + if ( + getattr(config, "sliding_window", None) is not None + and config.sliding_window < max_seq_len + and not any(windows) + ): + raise ValueError( + "sliding_window is shorter than context but no sliding layer was detected." + ) + if any(query_len > min(w, max_seq_len) for w in windows if w is not None): + raise ValueError("query_len exceeds the smallest sliding cache length.") + args, kwargs = exported_program.example_inputs + flat, spec = pytree.tree_flatten((args, kwargs)) + user = [ + s.arg.name + for s in exported_program.graph_signature.input_specs + if s.kind.name == "USER_INPUT" + ] + if len(user) != len(flat): + raise ValueError("The program's example inputs don't match its user inputs.") + lengths = [max_seq_len if w is None else min(w, max_seq_len) for w in windows] + gm = exported_program.module() + attention_layers = _cache_attention(gm.graph) + if len(attention_layers) != len(windows): + raise ValueError( + f"Expected {len(windows)} cache attention layers, matched {len(attention_layers)}." + ) + masks = {} + window_masks = [] + for i, (sdpa, _, _) in enumerate(attention_layers): + mask = sdpa.args[3] if len(sdpa.args) > 3 else sdpa.kwargs.get("attn_mask") + if mask is None or mask.op != "placeholder": + raise ValueError("Expected attention to use a mask input directly.") + if mask.name in masks and masks[mask.name] != lengths[i]: + raise ValueError( + "Layers with different cache lengths require separate masks." + ) + masks[mask.name] = lengths[i] + if windows[i] is not None: + window_masks.append(mask.name) + static = [] + state_index = 0 + for name, tensor in zip(user, flat): + if name == "input_ids": + static.append(torch.zeros(1, query_len, dtype=tensor.dtype)) + elif name == "position_ids": + static.append(torch.arange(query_len, dtype=tensor.dtype)[None]) + elif name in masks: + static.append( + torch.ones(1, 1, query_len, masks[name], dtype=torch.bool).tril() + ) + elif name.startswith(CACHE_PREFIX): + if tensor.dim() == 4: + static.append( + torch.zeros( + *tensor.shape[:2], + lengths[state_index // 3], + tensor.shape[3], + dtype=tensor.dtype, + ) + ) + else: + static.append(torch.zeros((), dtype=tensor.dtype)) + state_index += 1 + else: + raise ValueError(f"Unexpected hf_static input {name!r}.") + if state_index != 3 * len(windows): + raise ValueError("Expected three cache state inputs per layer.") + args, kwargs = pytree.tree_unflatten(static, spec) + program = torch.export.export(gm, args, kwargs, strict=False) + program.graph_module.meta.update( + hf_static_max_seq_len=max_seq_len, + hf_static_windows=windows, + hf_static_window_masks=list(dict.fromkeys(window_masks)), + ) + return program + + +def _cpu_call(graph): + def call(target, args, kwargs=None): + node = graph.call_function(target, args, kwargs) + node.meta.setdefault("custom", {})[CPU_META] = True + return node + + return call + + +def _masked( + call, scores: torch.fx.Node, keep: torch.fx.Node, masked_softmax: bool +) -> torch.fx.Node: + if masked_softmax: + # HTP fuses this into MaskedSoftmax; it holds for any score range once + # quantized, where an additive mask would set the scores' encoding. + low = call( + aten.add.Tensor, (call(aten.amin.default, (scores, [-1], True)), -20.0) + ) + return call(aten.where.self, (keep, scores, low)) + bias = call( + aten.where.self, + ( + keep, + call(aten.zeros_like.default, (scores,)), + call(aten.full_like.default, (scores, _MASK_FILL)), + ), + ) + return call(aten.add.Tensor, (scores, bias)) + + +def _constant( + gm: torch.fx.GraphModule, name: str, value: torch.Tensor +) -> torch.fx.Node: + if not hasattr(gm, name): + gm.register_buffer(name, value) + return gm.graph.get_attr(name) + + +def _step_mask(gm, mask, index, start, length: int, steps: int) -> torch.fx.Node: + """Visibility over incoming cache slots followed by the step's new rows.""" + call = _cpu_call(gm.graph) + slots = _constant( + gm, f"_hf_static_slots_{length}", torch.arange(length).view(1, 1, 1, length) + ) + end = call(aten.add.Tensor, (start, steps)) + unwritten = call( + aten.logical_or.default, + (call(aten.lt.Tensor, (slots, start)), call(aten.ge.Tensor, (slots, end))), + ) + past = call(aten.logical_and.default, (mask, unwritten)) + new = call(aten.index_select.default, (mask, 3, index)) + return call(aten.cat.default, ([past, new], -1)) + + +def _window_mask(gm, start, window: int, steps: int) -> torch.fx.Node: + """Which of [a window-sized cache's slots, the step's new rows] each query + row sees, as HF's sliding mask lets it: the cache holds the last + ``window`` positions before ``start`` in order (see ``_sliding_writes``), + and a query sees the positions less than ``window`` behind it.""" + call = _cpu_call(gm.graph) + slots = _constant( + gm, f"_hf_static_slots_{window}", torch.arange(window).view(1, 1, 1, window) + ) + rows = _constant( + gm, f"_hf_static_rows_{steps}", torch.arange(steps).view(1, 1, steps, 1) + ) + cols = _constant( + gm, f"_hf_static_cols_{steps}", torch.arange(steps).view(1, 1, 1, steps) + ) + first = call(aten.clamp.default, (call(aten.sub.Tensor, (start, window)), 0)) + # Query row i is at start + i; slot j holds position first + j. + behind = call( + aten.sub.Tensor, + (call(aten.add.Tensor, (rows, start)), call(aten.add.Tensor, (slots, first))), + ) + past = call( + aten.logical_and.default, + (call(aten.lt.Tensor, (slots, start)), call(aten.lt.Scalar, (behind, window))), + ) + new = call( + aten.logical_and.default, + ( + call(aten.le.Tensor, (cols, rows)), + call(aten.lt.Scalar, (call(aten.sub.Tensor, (rows, cols)), window)), + ), + ) + return call(aten.cat.default, ([past, new], -1)) + + +def _sliding_writes(gm, cache, rows, start, valid, window: int, steps: int): + """Shift to keep the latest window, then write real rows in place. + + Padding writes repeat the final real row's index AND data, so colliding + indices store identical values irrespective of kernel write order. + """ + assert 0 < steps <= window + call = _cpu_call(gm.graph) + slots = _constant(gm, f"_hf_static_slot_index_{window}", torch.arange(window)) + steps_ = _constant(gm, f"_hf_static_steps_{steps}", torch.arange(steps)) + start, valid = (call(aten.view.default, (t, [1])) for t in (start, valid)) + first = call(aten.clamp.default, (call(aten.sub.Tensor, (start, window)), 0)) + end = call(aten.add.Tensor, (start, valid)) + first_after = call(aten.clamp.default, (call(aten.sub.Tensor, (end, window)), 0)) + shift = call(aten.sub.Tensor, (first_after, first)) + source = call( + aten.clamp.default, (call(aten.add.Tensor, (slots, shift)), None, window - 1) + ) + call( + aten.index_copy_.default, + (cache, 2, slots, call(aten.index_select.default, (cache, 2, source))), + ) + row = call(aten.minimum.default, (steps_, call(aten.sub.Tensor, (valid, 1)))) + slot = call(aten.sub.Tensor, (call(aten.add.Tensor, (row, start)), first_after)) + return call( + aten.index_copy_.default, + (cache, 2, slot, call(aten.index_select.default, (rows, 2, row))), + ) + + +def _attention_concat( + call, q, cache_k, cache_v, k_new, v_new, keep, scale, masked_softmax +): + """kv_cache_io's shape of attention: one QK and one PV matmul over the cache + and the new rows, grouped-query heads by expanding K and V, which QNN splits + per head (mha2sha).""" + _, hq, steps, d = q.meta["val"].shape + _, hkv, length, _ = cache_k.meta["val"].shape + rep, total = hq // hkv, length + steps + k_all = call( + aten.cat.default, + ( + [ + call(aten.transpose.int, (cache_k, 2, 3)), + call(aten.transpose.int, (k_new, 2, 3)), + ], + 3, + ), + ) + v_all = call(aten.cat.default, ([cache_v, v_new], 2)) + if rep > 1: + + def expand(t, shape): + t = call( + aten.expand.default, + (call(aten.unsqueeze.default, (t, 2)), [1, hkv, rep, *shape]), + ) + return call(aten.reshape.default, (t, [1, hq, *shape])) + + k_all, v_all = expand(k_all, [d, total]), expand(v_all, [total, d]) + scores = call(aten.mul.Tensor, (call(aten.matmul.default, (q, k_all)), scale)) + probs = call( + aten._softmax.default, (_masked(call, scores, keep, masked_softmax), -1, False) + ) + return call(aten.matmul.default, (probs, v_all)) + + +def _attention_split( + call, q, cache_k, cache_v, k_new, v_new, keep, scale, masked_softmax +): + """The cache and the new rows in separate matmuls, grouped-query heads folded + into the query rows, so each KV head's cache is read once and never copied: + + scores = [q @ K_cache^T, q @ K_new^T] + out = softmax(scores)[:L] @ V_cache + softmax(scores)[L:] @ V_new + + Quantized, the two probability slices and partial products are calibrated + separately, so this assumes min/max observers (as the LLM recipes use): + how attention divides between the cache and the new rows varies from step + to step, and a moving average would clip the steps it doesn't favor. + + ``keep`` is the step mask repeated for each query head of a KV group.""" + _, hq, steps, d = q.meta["val"].shape + _, hkv, length, _ = cache_k.meta["val"].shape + rep = hq // hkv + qg = call(aten.reshape.default, (q, [1, hkv, rep * steps, d])) + scores = call( + aten.cat.default, + ( + [ + call( + aten.matmul.default, (qg, call(aten.transpose.int, (cache_k, 2, 3))) + ), + call( + aten.matmul.default, (qg, call(aten.transpose.int, (k_new, 2, 3))) + ), + ], + -1, + ), + ) + scores = call(aten.mul.Tensor, (scores, scale)) + probs = call( + aten._softmax.default, (_masked(call, scores, keep, masked_softmax), -1, False) + ) + p_past = call(aten.slice.Tensor, (probs, -1, 0, length)) + p_new = call(aten.slice.Tensor, (probs, -1, length, length + steps)) + out = call( + aten.add.Tensor, + ( + call(aten.matmul.default, (p_past, cache_v)), + call(aten.matmul.default, (p_new, v_new)), + ), + ) + return call(aten.reshape.default, (out, [1, hq, steps, d])) + + +def _cache_attention(graph: torch.fx.Graph): + """Each attention reading a cache the step writes, with its K and V writes.""" + layers = [] + for sdpa in graph.nodes: + if ( + sdpa.op != "call_function" + or sdpa.target != aten.scaled_dot_product_attention.default + ): + continue + k_write = cache_graph.producer_cache_write(sdpa.args[1]) + v_write = cache_graph.producer_cache_write(sdpa.args[2]) + if k_write is not None and v_write is not None: + layers.append((sdpa, k_write, v_write)) + if not layers: + raise ValueError("No attention reading a written cache was found.") + return layers + + +def _fold_small_constants(gm: torch.fx.GraphModule, max_numel: int = 1 << 16) -> None: + """Evaluate small values computed from constants alone (e.g. RoPE's + ``inv_freq`` cast and reshape) into constants, so that the position tables + take in all of RoPE: otherwise its angle is computed on HTP, which loses + precision once quantized. Large values (weights) are left alone.""" + graph = gm.graph + count = 0 + for node in list(graph.nodes): + if node.op != "call_function" or not node.all_input_nodes: + continue + if any(i.op != "get_attr" for i in node.all_input_nodes): + continue + val = node.meta.get("val") + if not isinstance(val, torch.Tensor) or val.numel() > max_numel: + continue + + def value(a): + if isinstance(a, torch.fx.Node): + obj = gm + for atom in a.target.split("."): + obj = getattr(obj, atom) + return obj + return a + + with torch.no_grad(): + folded = node.target( + *torch.fx.node.map_arg(node.args, value), + **torch.fx.node.map_arg(node.kwargs, value), + ) + if not isinstance(folded, torch.Tensor): + continue + with graph.inserting_before(node): + constant = _constant( + gm, f"_hf_static_constant_{count}", folded.contiguous() + ) + count += 1 + constant.meta = dict(node.meta) + node.replace_all_uses_with(constant) + graph.erase_node(node) + + +def _name_tables_by_shape(gm: torch.fx.GraphModule) -> None: + """Methods lowered together share constants by name; the position tables of + a prompt chunk and of a single token can differ in shape under the same name.""" + renamed = set() + for node in gm.graph.nodes: + if node.op != "get_attr" or not node.target.startswith("_qnn_position_table_"): + continue + table = getattr(gm, node.target) + name = f"{node.target}_{'x'.join(str(d) for d in table.shape)}" + if not hasattr(gm, name): + gm.register_buffer(name, table) + renamed.add(node.target) + node.target = name + for name in renamed: + del gm._buffers[name] + + +def _rewrite_attention( # noqa: C901 + gm: torch.fx.GraphModule, + masked_softmax: bool, + form: str, + windows: Sequence[Optional[int]], +) -> int: + """Attend over incoming caches and new rows, then update caches on CPU.""" + graph = gm.graph + call, cpu = graph.call_function, _cpu_call(graph) + layers = _cache_attention(graph) + if len(layers) != len(windows): + raise ValueError( + f"Expected {len(windows)} cache attention layers, matched {len(layers)}." + ) + attention = {"concat": _attention_concat, "split": _attention_split}[form] + masks, windowed = {}, {} + for i, (sdpa, k_write, v_write) in enumerate(layers): + q = sdpa.args[0] + mask = sdpa.args[3] if len(sdpa.args) > 3 else sdpa.kwargs.get("attn_mask") + if mask is None or mask.meta["val"].dtype != torch.bool: + raise ValueError(f"{sdpa.name}: expected a bool attention mask.") + cache_k, cache_v = k_write.args[0], v_write.args[0] + length, steps = cache_k.meta["val"].shape[2], q.meta["val"].shape[2] + rep = q.meta["val"].shape[1] // cache_k.meta["val"].shape[1] + sliding = windows[i] is not None + mask.meta.setdefault("custom", {})[CPU_META] = True + for write in (k_write, v_write): + write.meta.setdefault("custom", {})[CPU_META] = True + stack = [write.args[2]] + seen = set() + while stack: + node = stack.pop() + if node in seen: + continue + seen.add(node) + node.meta.setdefault("custom", {})[CPU_META] = True + stack.extend(node.all_input_nodes) + with graph.inserting_before(sdpa): + if mask not in masks: + index = k_write.args[2] + start = cpu( + aten.view.default, + (cpu(aten.slice.Tensor, (index, 0, 0, 1)), [1, 1, 1, 1]), + ) + keep = ( + _window_mask(gm, start, length, steps) + if sliding + else _step_mask(gm, mask, index, start, length, steps) + ) + masks[mask] = (keep, start) + keep, start = masks[mask] + if form == "split" and rep > 1: + if (mask, rep) not in masks: + masks[(mask, rep)] = cpu( + aten.repeat.default, (keep, [1, 1, rep, 1]) + ) + keep = masks[(mask, rep)] + scale = sdpa.args[6] if len(sdpa.args) > 6 else sdpa.kwargs.get("scale") + out = attention( + call, + q, + cache_k, + cache_v, + k_write.args[3], + v_write.args[3], + keep, + scale or q.meta["val"].shape[3] ** -0.5, + masked_softmax, + ) + out.meta = dict(sdpa.meta) + sdpa.replace_all_uses_with(out) + graph.erase_node(sdpa) + if sliding: + windowed[k_write] = windowed[v_write] = (mask, start, length, steps) + graph.eliminate_dead_code() + output = next(n for n in graph.nodes if n.op == "output") + valid = {} + for _, k_write, v_write in layers: + for write in (k_write, v_write): + if write not in windowed: + output.prepend(write) + continue + mask, start, length, steps = windowed[write] + with graph.inserting_before(output): + if mask not in valid: + rows = cpu(aten.any.dim, (mask, -1, True)) + valid[mask] = cpu(aten.sum.dim_IntList, (rows, [2], True)) + shifted = _sliding_writes( + gm, write.args[0], write.args[3], start, valid[mask], length, steps + ) + write.replace_all_uses_with(shifted) + graph.erase_node(write) + # Counter updates are part of the mutable-input protocol, not model compute. + for node in graph.nodes: + if ( + node.op == "call_function" + and node.target == aten.add_.Tensor + and node.args[0].meta.get("custom", {}).get(CPU_META) + ): + node.meta.setdefault("custom", {})[CPU_META] = True + return len(layers) + + +def to_hf_static( + exported_program: ExportedProgram, + masked_softmax: bool = False, + attention: str = "split", +) -> ExportedProgram: + """Rewrite a specialized graph without changing its per-step signature.""" + if attention not in ("split", "concat"): + raise ValueError(f"Unknown attention form {attention!r}.") + gm = exported_program.module() + graph = gm.graph + position_ids = next( + n for n in graph.nodes if n.op == "placeholder" and n.name == "position_ids" + ) + layout = HfStaticLayout.of(exported_program) + _fold_small_constants(gm) + # Only RoPE depends on the position alone; the cache writes follow the counters. + position_ids.meta.setdefault("custom", {})[CPU_META] = True + cache_graph.fold_position_tables( + gm, + position_ids, + position_ids, + layout.max_seq_len, + layout.dtype, + index_meta={"custom": {CPU_META: True}}, + ) + cache_graph.split_rotate_half(graph) + cache_graph.fold_weight_arithmetic(gm) + cache_graph.norm_after_head_transpose(graph) + # Unlifting the program later (calibration, lowering) fails on folded + # values registered as non-persistent buffers. + gm._non_persistent_buffers_set.clear() + _name_tables_by_shape(gm) + windows = exported_program.graph_module.meta.get( + "hf_static_windows", [None] * layout.n_layers + ) + matched = _rewrite_attention(gm, masked_softmax, attention, windows) + if matched != layout.n_layers: + raise ValueError( + f"Expected {layout.n_layers} cache attention layers, matched {matched}." + ) + graph.eliminate_dead_code() + graph.lint() + gm.recompile() + args, kwargs = copy.deepcopy(exported_program.example_inputs) + with torch.fx.traceback.preserve_node_meta(): + program = torch.export.export(gm, args, kwargs, strict=False) + program.graph_module.meta.update(exported_program.graph_module.meta) + return program + + +def tag_cache_mutations(program: ExportedProgram) -> None: + """Tag copies introduced by functionalization of the cache protocol.""" + state_names = [ + s.arg.name + for s in program.graph_signature.input_specs + if s.kind.name == "USER_INPUT" and s.arg.name.startswith(CACHE_PREFIX) + ] + counters = set(state_names[2::3]) + nodes = {n.name: n for n in program.graph.nodes} + outputs = [ + (spec, spec.target) + for spec in program.graph_signature.output_specs + if spec.kind.name == "USER_INPUT_MUTATION" and spec.target in state_names + ] + user_outputs = [ + spec + for spec in program.graph_signature.output_specs + if spec.kind.name == "USER_OUTPUT" + ] + if len(user_outputs) == len(state_names) + 1: + outputs.extend(zip(user_outputs[1:], state_names)) + for spec, target in outputs: + stack = [nodes[spec.arg.name]] + seen = set() + while stack: + node = stack.pop() + if node in seen: + continue + seen.add(node) + if node.op != "call_function": + continue + try: + op = unwrap_op_overload(node.target) + except TypeError: + continue + if target in counters: + node.meta.setdefault("custom", {})[CPU_META] = True + stack.extend(node.all_input_nodes) + elif op in (aten.copy.default, aten.copy_.default): + node.meta.setdefault("custom", {})[CPU_META] = True + stack.append(node.args[1]) + elif op in ( + aten.index_put.default, + aten.index_put_.default, + aten.index_copy.default, + aten.index_copy_.default, + ): + node.meta.setdefault("custom", {})[CPU_META] = True + stack.append(node.args[0]) + + +def flat_call( + exported_program: ExportedProgram, module: Optional[torch.nn.Module] = None +): + """The program (or ``module``, e.g. its PT2E-prepared form) as a function of + its flat user inputs, returning its flat outputs.""" + module = module if module is not None else exported_program.module() + spec = exported_program.call_spec.in_spec + + def call(*flat): + args, kwargs = pytree.tree_unflatten(list(flat), spec) + return pytree.tree_flatten(module(*args, **kwargs))[0] + + return call diff --git a/extension/llm/export/test/BUCK b/extension/llm/export/test/BUCK index 9a40024abbe..c1efc134289 100644 --- a/extension/llm/export/test/BUCK +++ b/extension/llm/export/test/BUCK @@ -9,6 +9,17 @@ load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") oncall("executorch") +fbcode_target(_kind = runtime.python_test, + name = "test_hf_static", + srcs = ["test_hf_static.py"], + deps = [ + "//caffe2:torch", + "//executorch/extension/llm/export:hf_static", + "//executorch/exir:lib", + "//executorch/runtime:runtime", + ], +) + fbcode_target(_kind = runtime.python_test, name = "test_builder", srcs = ["test_builder.py"], diff --git a/extension/llm/export/test/test_hf_static.py b/extension/llm/export/test/test_hf_static.py new file mode 100644 index 00000000000..9cd56fb00ba --- /dev/null +++ b/extension/llm/export/test/test_hf_static.py @@ -0,0 +1,141 @@ +# 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. + +"""Graph-level regressions for hf_static configuration and padding writes.""" + +import unittest +from types import SimpleNamespace + +import torch +from executorch.exir import EdgeCompileConfig, ExecutorchBackendConfig, to_edge +from executorch.exir.passes import MemoryPlanningPass +from executorch.extension.llm.export import hf_static +from executorch.extension.llm.export.cache_graph import fold_position_tables +from executorch.runtime import Runtime, Verification + + +class TestHfStaticGraph(unittest.TestCase): + def _writes(self, valid, start=8): + graph = torch.fx.Graph() + cache, rows, position, count = [ + graph.placeholder(n) for n in ("cache", "rows", "position", "count") + ] + graph.output(cache) + gm = torch.fx.GraphModule(torch.nn.Module(), graph) + graph.erase_node(next(n for n in graph.nodes if n.op == "output")) + out = hf_static._sliding_writes(gm, cache, rows, position, count, 6, 4) + graph.output(out) + gm.recompile() + payloads = [] + + class Observe(torch.fx.Interpreter): + def run_node(self, node): + if node.target == torch.ops.aten.index_copy_.default: + args, _ = self.fetch_args_kwargs_from_env(node) + payloads.append((args[2].clone(), args[3].clone())) + return super().run_node(node) + + cache = torch.arange(12.0).reshape(1, 1, 6, 2) + rows = torch.arange(100.0, 108.0).reshape(1, 1, 4, 2) + result = Observe(gm).run(cache, rows, torch.tensor(start), torch.tensor(valid)) + return result, payloads + + def test_reinplace_write_chain_has_no_full_cache_copy(self): + class Chain(torch.nn.Module): + def forward(self, cache, rows, index): + cache.index_copy_(2, index, rows) + checksum = cache.sum() + cache.index_copy_(2, index, rows + 1) + return checksum, cache + + inputs = (torch.zeros(1, 1, 6, 2), torch.ones(1, 1, 2, 2), torch.tensor([1, 2])) + program = torch.export.export(Chain(), inputs) + manager = to_edge( + program, compile_config=EdgeCompileConfig(_check_ir_validity=False) + ).to_executorch( + ExecutorchBackendConfig( + run_reinplace_pass=True, + memory_planning_pass=MemoryPlanningPass(alloc_graph_input=False), + ) + ) + plan = manager.executorch_program.execution_plan[0] + ops = [op.name for op in plan.operators] + self.assertNotIn("aten::index_put", ops) + self.assertNotIn("aten::copy_", ops) + self.assertEqual(plan.outputs[0], plan.inputs[0]) + runtime = Runtime.get().load_program( + manager.buffer, verification=Verification.Minimal + ) + outputs = runtime.load_method("forward").execute(inputs) + self.assertEqual(float(outputs[1]), 4.0) + torch.testing.assert_close( + inputs[0][..., 1:3, :], torch.full((1, 1, 2, 2), 2.0) + ) + + def test_full_context_position_table_probe(self): + class Position(torch.nn.Module): + def forward(self, positions, values): + return values + torch.cos(positions.float())[..., None] + + inputs = (torch.arange(4).reshape(1, 4), torch.zeros(1, 4, 2)) + program = torch.export.export(Position(), inputs) + gm = program.module() + position = next( + n for n in gm.graph.nodes if n.op == "placeholder" and n.name == "positions" + ) + fold_position_tables(gm, position, position, 4, torch.float32) + gm.graph.eliminate_dead_code() + gm.recompile() + torch.testing.assert_close(gm(*inputs), Position()(*inputs)) + + def test_padding_collision_payloads_are_identical(self): + for valid in (1, 2, 3, 4): + with self.subTest(valid=valid): + result, payloads = self._writes(valid) + slots, rows = payloads[-1] + for i in range(4): + for j in range(i + 1, 4): + if slots[i] == slots[j]: + torch.testing.assert_close( + rows[:, :, i], rows[:, :, j], rtol=0, atol=0 + ) + expected = torch.cat( + ( + torch.arange(12.0).reshape(1, 1, 6, 2)[:, :, valid:], + torch.arange(100.0, 108.0).reshape(1, 1, 4, 2)[:, :, :valid], + ), + 2, + ) + torch.testing.assert_close(result, expected) + + def test_prefill_longer_than_window_is_rejected(self): + config = SimpleNamespace(num_hidden_layers=2, sliding_window=6) + # The graph is deliberately absent: validation must precede tracing. + with self.assertRaisesRegex(ValueError, "smallest sliding cache"): + hf_static.specialize(None, 8, 32, config=config) + + def test_all_sliding_config_without_layer_types(self): + config = SimpleNamespace(num_hidden_layers=2, sliding_window=6) + self.assertEqual(hf_static.sliding_windows(config), [6, 6]) + config.layer_types = ["sliding_attention", "full_attention"] + self.assertEqual(hf_static.sliding_windows(config), [6, None]) + + def test_inconsistent_window_config_is_rejected(self): + config = SimpleNamespace( + num_hidden_layers=2, + sliding_window=6, + layer_types=["full_attention", "full_attention"], + ) + with self.assertRaisesRegex(ValueError, "no sliding layer"): + hf_static.specialize(None, 4, 32, config=config) + + def test_oversized_write_is_rejected(self): + with self.assertRaises(AssertionError): + hf_static._sliding_writes(None, None, None, None, None, 3, 4) + + +if __name__ == "__main__": + unittest.main() diff --git a/extension/llm/io_contracts/BUCK b/extension/llm/io_contracts/BUCK new file mode 100644 index 00000000000..f2b61f1f32d --- /dev/null +++ b/extension/llm/io_contracts/BUCK @@ -0,0 +1,44 @@ +load("@fbcode_macros//build_defs:build_file_migration.bzl", "fbcode_target") +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") + +oncall("executorch") + +fbcode_target(_kind = runtime.python_library, + name = "io_contracts", + srcs = [ + "__init__.py", + "_transformers.py", + "hf_static.py", + "kv_cache_io.py", + ], + base_module = "executorch.extension.llm.io_contracts", + visibility = ["PUBLIC"], + deps = [ + "//caffe2:torch", + "//executorch/extension/llm/export:hf_static", + ], +) + +fbcode_target(_kind = runtime.python_test, + name = "test_kv_cache_io", + srcs = ["test/test_kv_cache_io.py"], + deps = [ + ":io_contracts", + "//caffe2:torch", + "//executorch/extension/llm/export:hf_static", + "//executorch/exir:lib", + "//executorch/runtime:runtime", + ], +) + +fbcode_target(_kind = runtime.python_test, + name = "test_hf_static", + srcs = ["test/test_hf_static.py"], + deps = [ + ":io_contracts", + "//caffe2:torch", + "//executorch/extension/llm/export:hf_static", + "//executorch/exir:lib", + "//executorch/runtime:runtime", + ], +) diff --git a/extension/llm/io_contracts/__init__.py b/extension/llm/io_contracts/__init__.py new file mode 100644 index 00000000000..cc2b4ffccc7 --- /dev/null +++ b/extension/llm/io_contracts/__init__.py @@ -0,0 +1,31 @@ +# 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. + +"""Graph I/O contracts ExecuTorch programs may follow, with their Python drivers. + +Importing this package registers each contract's runner with Hugging Face +``transformers`` exporters when it is installed, so ``ExportedGenerator`` drives +programs labelled with the contract. +""" + +from executorch.extension.llm.io_contracts.hf_static import ( # noqa: F401 + HfStaticDriver, + HfStaticRunner, +) +from executorch.extension.llm.io_contracts.kv_cache_io import ( # noqa: F401 + CONTRACT as KV_CACHE_IO, + io_contract, + KvCacheIORunner, +) + +try: + from transformers.exporters import register_io_contract +except ImportError: + pass +else: + from executorch.extension.llm.io_contracts._transformers import KvCacheIOModelRunner + + register_io_contract(KV_CACHE_IO, KvCacheIOModelRunner) diff --git a/extension/llm/io_contracts/_transformers.py b/extension/llm/io_contracts/_transformers.py new file mode 100644 index 00000000000..053ff33b97e --- /dev/null +++ b/extension/llm/io_contracts/_transformers.py @@ -0,0 +1,167 @@ +# 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. + +"""``transformers`` exporters runner for ``kvcache_io`` programs.""" + +import torch +from executorch.extension.llm.io_contracts.hf_static import HfStaticRunner +from executorch.extension.llm.io_contracts.kv_cache_io import KvCacheIORunner +from transformers.exporters import ModelRunner +from transformers.exporters.metadata import ExportMetadata + + +class KvCacheIOModelRunner(ModelRunner): + """Runs a ``kvcache_io`` program under ``ExportedGenerator``. + + The program keeps the cache in buffers this runner owns, so it is a + state-owning runner: the generation loop feeds token ids and reads logits, + and a prompt program hands its filled cache to the decode program through + ``state_tensors`` / ``adopt_state``. + """ + + owns_state = True + + def __init__(self, program, export_metadata=None, **kwargs): + self.driver = KvCacheIORunner([program]) + self.export_metadata = ExportMetadata.from_dict(export_metadata) + # One program with both steps serves the whole loop. + self.drives_generation = ( + self.driver.prompt_method is not self.driver.step_method + ) + layers = self.driver.num_layers + self.state_paths = frozenset( + [f"past_k_{i}" for i in range(layers)] + + [f"past_v_{i}" for i in range(layers)] + ) + self.input_names = ("input_ids",) + self.cache_inputs = ("past_key_values",) + if self.export_metadata.dtype is None: + self.dtype = torch.float32 + + @property + def state_length(self) -> int: + return self.driver.filled + + def state_tensors(self, paths=None) -> dict: + caches = self.driver.past() + layers = self.driver.num_layers + named = {f"past_k_{i}": caches[i] for i in range(layers)} + named.update({f"past_v_{i}": caches[layers + i] for i in range(layers)}) + return {k: v for k, v in named.items() if paths is None or k in paths} + + def adopt_state(self, leaves: dict, length: int) -> None: + layers = self.driver.num_layers + self.driver.adopt( + [leaves[f"past_k_{i}"] for i in range(layers)] + + [leaves[f"past_v_{i}"] for i in range(layers)], + length, + ) + + def reset_state(self) -> None: + self.driver.reset() + + def __call__(self, input_ids=None, **kwargs) -> dict: + if input_ids is None or input_ids.shape[0] != 1: + raise ValueError("kvcache_io programs take one sequence of input_ids.") + logits = self.driver.forward(input_ids[0].tolist()) + return {"logits": logits.view(1, 1, -1)} + + +class HfStaticModelRunner(ModelRunner): + """State-owning HF generator adapter for a fixed prefill/decode program.""" + + owns_state = True + drives_generation = True + + def __init__(self, program, export_metadata=None, **kwargs): + self.driver = HfStaticRunner(program) + metadata = dict(export_metadata or {}) + metadata["kwargs"] = { + **metadata.get("kwargs", {}), + "attention_mask": {"rank": 2, "dtype": "bool"}, + } + self.export_metadata = ExportMetadata.from_dict(metadata) + self.input_names = ("input_ids", "attention_mask", "position_ids") + self.cache_inputs = ("past_key_values",) + self.state_paths = frozenset() + self.dtype = self.driver.driver.layout.dtype + + @property + def state_length(self): + return self.driver.filled + + def reset_state(self): + self.driver.reset() + + def _validate_mask(self, mask, query_len): + message = "hf_static does not support caller-padded or custom attention masks." + if ( + not isinstance(mask, torch.Tensor) + or mask.ndim not in (2, 4) + or mask.shape[0] != 1 + ): + raise ValueError(message) + start = self.driver.filled + if mask.ndim == 2: + width = min(mask.shape[1], start + query_len) + if mask.shape[1] < query_len or not bool((mask[:, :width] == 1).all()): + raise ValueError(message) + return + if tuple(mask.shape[:3]) != (1, 1, query_len) or mask.shape[-1] < query_len: + raise ValueError(message) + if mask.dtype == torch.bool: + keep = mask + elif mask.dtype.is_floating_point and bool( + ((mask == 0) | (mask <= -10000)).all() + ): + keep = mask == 0 + else: + raise ValueError(message) + width = mask.shape[-1] + first = max(0, start + query_len - width) + slots = torch.arange(first, first + width, device=mask.device) + positions = torch.arange(start, start + query_len, device=mask.device) + causal = (slots[None] <= positions[:, None])[None, None] + if torch.equal(keep, causal): + return + windows = set(self.driver.driver.layout.mask_widths.values()) + for window in windows: + windowed = ( + causal & ((positions[:, None] - slots[None]) < window)[None, None] + ) + if torch.equal(keep, windowed): + return + raise ValueError(message) + + def __call__( + self, input_ids=None, attention_mask=None, position_ids=None, **kwargs + ): + if input_ids is None or input_ids.ndim != 2 or input_ids.shape[0] != 1: + raise ValueError("hf_static programs take one sequence of input_ids.") + if attention_mask is not None: + masks = ( + attention_mask.values() + if isinstance(attention_mask, dict) + else (attention_mask,) + ) + for mask in masks: + if mask is None: + continue + self._validate_mask(mask, input_ids.shape[1]) + if position_ids is not None: + expected = torch.arange( + self.driver.filled, + self.driver.filled + input_ids.shape[1], + device=position_ids.device, + )[None] + if position_ids.shape != expected.shape or not torch.equal( + position_ids, expected + ): + raise ValueError( + "hf_static requires consecutive absolute position_ids." + ) + logits = self.driver.forward(input_ids[0].tolist()) + return {"logits": logits.view(1, 1, -1)} diff --git a/extension/llm/io_contracts/hf_static.py b/extension/llm/io_contracts/hf_static.py new file mode 100644 index 00000000000..d9b30e88f6c --- /dev/null +++ b/extension/llm/io_contracts/hf_static.py @@ -0,0 +1,333 @@ +# 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. + +"""Python drivers for the version-1 ``hf_static`` float-cache contract.""" + +from typing import Callable, Dict, List, Sequence + +import torch +from executorch.extension.llm.export.hf_static import ( + CACHE_PREFIX, + DECODE_METHOD, + HfStaticLayout, + IO_CONTRACT, + IO_CONTRACT_VERSION, + PREFILL_METHOD, +) +from executorch.extension.llm.io_contracts.kv_cache_io import ( + constant_method, + tensor_dtype, +) + + +class HfStaticDriver: + """Drive flat static-cache steps for calibration and eager validation. + + ``steps`` maps each query length to a callable returning logits followed by + layer-ordered keys, values and counters. With ``in_place``, keep the caller's + buffers instead of adopting the returned state. + """ + + def __init__( + self, steps: Dict[int, Callable], layout: HfStaticLayout, in_place: bool = False + ): + self.steps, self.layout, self.in_place = steps, layout, in_place + self.state = [ + torch.zeros(shape, dtype=dtype) + for shape, dtype in zip(layout.state_shapes, layout.state_dtypes) + ] + self.filled = 0 + + def reset(self) -> None: + for tensor in self.state: + tensor.zero_() + self.filled = 0 + + def forward( # noqa: C901 + self, tokens: Sequence[int], start: int, query_len: int + ) -> torch.Tensor: + """Run a nonempty chunk at ``start`` and return its last real logits.""" + n, length = len(tokens), self.layout.max_seq_len + if not 0 < n <= query_len or query_len not in self.steps: + raise ValueError("The chunk must fit an available query length.") + if start < 0 or start + query_len > length: + raise ValueError( + f"Query at {start} with length {query_len} exceeds context {length}." + ) + pos = torch.arange(start, start + query_len, dtype=self.layout.position_dtype) + inputs, state = [], iter(self.state) + for tensor in self.state: + if tensor.dim() == 0: + tensor.fill_(start) + for name in self.layout.input_names: + if name == "input_ids": + ids = list(tokens) + [tokens[-1]] * (query_len - n) + inputs.append(torch.tensor([ids], dtype=self.layout.token_dtype)) + elif name == "position_ids": + inputs.append(pos[None]) + elif name.startswith("attention_mask"): + width = self.layout.mask_widths[name] + mask = torch.arange(width)[None] <= pos[:, None] + mask[n:] = False + inputs.append(mask[None, None]) + else: + inputs.append(next(state)) + out = self.steps[query_len](*inputs) + if not self.in_place: + self.state = [t.clone() for t in out[1:]] + for tensor in self.state: + if tensor.dim() == 0: + tensor.fill_(start + n) + self.filled = start + n + return out[0].reshape(-1, out[0].shape[-1])[n - 1] + + def run(self, tokens: Sequence[int]) -> torch.Tensor: + """Append tokens, using decode for a tail whose padded chunk would overflow.""" + tokens = list(tokens) + if not tokens or self.filled + len(tokens) > self.layout.max_seq_len: + raise ValueError("Tokens must be nonempty and fit the remaining context.") + offset, chunk, logits = 0, max(self.steps), None + while offset < len(tokens): + q = chunk + if len(tokens) - offset == 1 or self.filled + q > self.layout.max_seq_len: + q = 1 + if q not in self.steps: + raise ValueError( + "A single-token decode method is required for this tail." + ) + part = tokens[offset : offset + q] + logits = self.forward(part, self.filled, q) + offset += len(part) + return logits + + def generate(self, prompt: Sequence[int], max_new_tokens: int) -> List[int]: + """Greedy sequence including the prompt, for export-time calibration.""" + self.reset() + tokens = list(prompt) + if max_new_tokens < 0: + raise ValueError("max_new_tokens must be nonnegative.") + logits = self.run(tokens) + for _ in range(max_new_tokens): + if len(tokens) >= self.layout.max_seq_len: + break + tokens.append(int(logits.argmax())) + logits = self.run(tokens[-1:]) + return tokens + + +class HfStaticRunner: + """Own the shared cache buffers of a loaded ``hf_static`` program.""" + + def __init__(self, program): # noqa: C901 + version = constant_method(program, "get_io_contract_version") + if ( + not isinstance(version, int) + or isinstance(version, bool) + or not 1 <= version <= IO_CONTRACT_VERSION + ): + raise ValueError( + f"Unsupported hf_static contract version {version!r}; supported: 1." + ) + if constant_method(program, "get_io_contract") != IO_CONTRACT: + raise ValueError("Expected an hf_static program.") + length = constant_method(program, "get_max_seq_len") + layers = constant_method(program, "get_n_layers") + dtype_name = constant_method(program, "get_cache_dtype") + if not isinstance(length, int) or isinstance(length, bool) or length <= 0: + raise ValueError("get_max_seq_len must be a positive integer.") + if not isinstance(layers, int) or isinstance(layers, bool) or layers <= 0: + raise ValueError("get_n_layers must be a positive integer.") + self.methods = {} + layouts = {} + for name in (DECODE_METHOD, PREFILL_METHOD): + if name not in program.method_names: + continue + method = program.load_method(name) + meta = method.metadata + count = meta.num_inputs() + if count < 3 + 3 * layers: + raise ValueError(f"{name}: not enough inputs for {layers} layers.") + inputs = [meta.input_tensor_meta(i) for i in range(count)] + shapes = [torch.Size(t.sizes()) for t in inputs] + dtypes = [tensor_dtype(t) for t in inputs] + masks = count - 2 - 3 * layers + query = shapes[0][-1] if len(shapes[0]) == 2 else 0 + if query <= 0 or shapes[0] != (1, query) or dtypes[0] != torch.int64: + raise ValueError(f"{name}: invalid input_ids shape or dtype.") + if name == DECODE_METHOD and query != 1: + raise ValueError("decode must take exactly one token.") + if shapes[1 + masks] != (1, query) or dtypes[1 + masks] != torch.int64: + raise ValueError(f"{name}: invalid position_ids shape or dtype.") + widths = {} + for i in range(masks): + shape = shapes[1 + i] + if ( + len(shape) != 4 + or shape[:3] != (1, 1, query) + or shape[-1] <= 0 + or shape[-1] > length + or dtypes[1 + i] != torch.bool + ): + raise ValueError(f"{name}: invalid attention mask shape or dtype.") + widths[f"attention_mask_{i}"] = shape[-1] + state_shapes, state_dtypes = shapes[-3 * layers :], dtypes[-3 * layers :] + cache_shapes = [] + for i in range(layers): + k, v, counter = state_shapes[3 * i : 3 * i + 3] + kd, vd, cd = state_dtypes[3 * i : 3 * i + 3] + if ( + len(k) != 4 + or k != v + or k[0] != 1 + or any(d <= 0 for d in k) + or k[2] > length + or counter != () + or cd != torch.int64 + ): + raise ValueError(f"{name}: invalid cache shapes or counter.") + if ( + kd != vd + or kd not in (torch.float16, torch.float32) + or dtype_name != str(kd).removeprefix("torch.") + ): + raise ValueError( + f"{name}: get_cache_dtype does not match the float cache inputs." + ) + if k[2] not in widths.values() or query > k[2]: + raise ValueError( + f"{name}: mask/cache width mismatch or query longer than cache." + ) + cache_shapes.extend((k, v)) + if query > length or meta.num_outputs() not in ( + 1 + 3 * layers, + 1 + 6 * layers, + ): + raise ValueError(f"{name}: invalid context or output count.") + # ExecuTorch places mutated inputs before the user logits. + expected_state = sorted( + (tuple(shape), str(dtype)) + for shape, dtype in zip(state_shapes, state_dtypes) + ) + mutation_state = sorted( + ( + tuple(meta.output_tensor_meta(i).sizes()), + str(tensor_dtype(meta.output_tensor_meta(i))), + ) + for i in range(3 * layers) + ) + if mutation_state != expected_state: + raise ValueError( + f"{name}: mutated outputs differ from the cache inputs." + ) + if meta.num_outputs() == 1 + 6 * layers: + for i, (shape, dtype) in enumerate(zip(state_shapes, state_dtypes)): + info = meta.output_tensor_meta(3 * layers + 1 + i) + if torch.Size(info.sizes()) != shape or tensor_dtype(info) != dtype: + raise ValueError( + f"{name}: user state output {i} differs from its input." + ) + logits = meta.output_tensor_meta(3 * layers) + if ( + len(logits.sizes()) != 3 + or tuple(logits.sizes())[:2] != (1, query) + or not tensor_dtype(logits).is_floating_point + ): + raise ValueError(f"{name}: invalid logits shape or dtype.") + names = ["input_ids", *widths, "position_ids"] + [ + f"{CACHE_PREFIX}{i}" for i in range(3 * layers) + ] + layouts[name] = HfStaticLayout( + names, + query, + length, + layers, + cache_shapes, + state_shapes, + state_dtypes[0], + widths, + state_dtypes=state_dtypes, + token_dtype=dtypes[0], + position_dtype=dtypes[1 + masks], + window_masks=[n for n, w in widths.items() if w < length], + ) + self.methods[name] = method + if DECODE_METHOD not in self.methods: + raise ValueError("An hf_static program requires decode.") + layout = layouts[DECODE_METHOD] + steps = {} + for name, method in self.methods.items(): + other = layouts[name] + if ( + other.state_shapes, + other.state_dtypes, + other.mask_widths, + other.token_dtype, + other.position_dtype, + ) != ( + layout.state_shapes, + layout.state_dtypes, + layout.mask_widths, + layout.token_dtype, + layout.position_dtype, + ): + raise ValueError( + "prefill and decode must have identical cache and mask layouts." + ) + if other.query_len in steps: + raise ValueError("prefill and decode must have distinct query lengths.") + + def step(*inputs, method=method): + outputs = method.execute(inputs) + state = ( + outputs[3 * layers + 1 :] + if len(outputs) == 1 + 6 * layers + else inputs[-3 * layers :] + ) + return [outputs[3 * layers], *state] + + steps[other.query_len] = step + self.driver = HfStaticDriver(steps, layout, in_place=True) + self.eos_ids = constant_method(program, "get_eos_ids", []) + if not isinstance(self.eos_ids, list): + self.eos_ids = [self.eos_ids] + + @property + def filled(self): + return self.driver.filled + + def reset(self): + self.driver.reset() + + def forward(self, tokens): + return self.driver.run(tokens) + + def generate(self, prompt_ids, max_new_tokens): + """Return new greedy tokens only, stopping at EOS or a full context.""" + if max_new_tokens < 0: + raise ValueError("max_new_tokens must be nonnegative.") + self.reset() + logits = self.forward(prompt_ids) + generated = [] + while ( + len(generated) < max_new_tokens + and self.filled < self.driver.layout.max_seq_len + ): + token = int(logits.argmax()) + generated.append(token) + if token in self.eos_ids: + break + logits = self.forward([token]) + return generated + + +try: + from transformers.exporters import register_io_contract +except ImportError: + pass +else: + from executorch.extension.llm.io_contracts._transformers import HfStaticModelRunner + + register_io_contract(IO_CONTRACT, HfStaticModelRunner) diff --git a/extension/llm/io_contracts/kv_cache_io.py b/extension/llm/io_contracts/kv_cache_io.py new file mode 100644 index 00000000000..9e45d7964d5 --- /dev/null +++ b/extension/llm/io_contracts/kv_cache_io.py @@ -0,0 +1,244 @@ +# 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. + +"""Python driver for programs that speak the ``kvcache_io`` I/O contract. + +The contract (written by ``backends/qualcomm/utils/kv_cache_io.py``, driven in +C++ by ``KvCacheIOManager``) is one decoder step per method:: + + (input_tokens, atten_mask, pos_ids, past_k_0..N-1, past_v_0..N-1) + -> (logits, new_k_0..N-1, new_v_0..N-1) + +``kv_forward`` takes one token, an optional ``prefill_forward`` a fixed chunk. +Past K is stored transposed ``[1, H, D, P]`` and V as ``[1, H, P, D]``; the +step returns only its own K/V slice and reads the past under an attention mask +the caller writes. This driver owns those buffers, so callers only exchange +tokens and logits. +""" + +from typing import Dict, List, Optional, Sequence + +import torch + +CONTRACT = "kvcache_io" +PREFILL_METHOD = "prefill_forward" +DECODE_METHOD = "kv_forward" + + +# ExecuTorch ScalarType codes (c10 numbering) of the dtypes the contract uses. +_DTYPES = { + 0: torch.uint8, + 1: torch.int8, + 2: torch.int16, + 3: torch.int32, + 4: torch.int64, + 5: torch.float16, + 6: torch.float32, + 7: torch.float64, + 11: torch.bool, + 15: torch.bfloat16, + 27: torch.uint16, +} + + +def tensor_dtype(info) -> torch.dtype: + """A runtime tensor's declared dtype, accepting c10 codes or torch dtypes.""" + dtype = info.dtype() + if isinstance(dtype, torch.dtype): + return dtype + if int(dtype) not in _DTYPES: + raise ValueError(f"Unsupported tensor dtype code {dtype!r}.") + return _DTYPES[int(dtype)] + + +def constant_method(program, name: str, default=None): + if name not in program.method_names: + return default + value = program.load_method(name).execute([]) + return value[0] if len(value) == 1 else list(value) + + +def io_contract(program) -> Optional[str]: + """The I/O contract a loaded program names, or None.""" + return constant_method(program, "get_io_contract") + + +class _Method: + """One contract method with its own token, mask and position inputs.""" + + def __init__(self, program, name: str): + self.name = name + self.method = program.load_method(name) + meta = self.method.metadata + self.num_layers = (meta.num_inputs() - 3) // 2 + tokens, mask, positions = (meta.input_tensor_meta(i) for i in range(3)) + self.ar_len = tokens.sizes()[-1] + self.context_len = mask.sizes()[-1] + self.past_len = self.context_len - self.ar_len + self.tokens = torch.zeros(list(tokens.sizes()), dtype=tensor_dtype(tokens)) + self.mask = torch.zeros(list(mask.sizes()), dtype=tensor_dtype(mask)) + self.positions = torch.zeros( + list(positions.sizes()), dtype=tensor_dtype(positions) + ) + self.cache_shapes = [ + ( + list(meta.input_tensor_meta(i).sizes()), + tensor_dtype(meta.input_tensor_meta(i)), + ) + for i in range(3, 3 + 2 * self.num_layers) + ] + + def fill_mask(self, n_past: int, n_valid: int) -> None: + # Visible slots are 0 in a float mask and the top code in a quantized + # one; the program's encoding maps that code to 0. + dtype = self.mask.dtype + if dtype.is_floating_point: + hidden, visible = -65535.0, 0.0 + else: + hidden, visible = 0, torch.iinfo(dtype).max + self.mask.fill_(hidden) + rows = self.mask.view(self.ar_len, self.context_len) + rows[:, :n_past] = visible + for r in range(self.ar_len): + own = min(r, max(n_valid - 1, 0)) + rows[r, self.past_len : self.past_len + own + 1] = visible + + +class KvCacheIORunner: + """Greedy-capable driver for ``kvcache_io`` programs. + + ``programs`` holds one program with both methods (a hybrid export) or the + separate decode and prompt programs. Prompts run through + ``prefill_forward`` in chunks when there is one, else token by token. + """ + + def __init__(self, programs: Sequence): + self.methods: Dict[str, _Method] = {} + for program in programs: + for name in (DECODE_METHOD, PREFILL_METHOD): + if name in program.method_names and name not in self.methods: + self.methods[name] = _Method(program, name) + if not self.methods: + raise ValueError( + f"A {CONTRACT} program has a {DECODE_METHOD} or {PREFILL_METHOD} method." + ) + # Single tokens go through the decode step, prompts through the + # prompt step; a program with only one serves both. + self.step_method = ( + self.methods.get(DECODE_METHOD) or self.methods[PREFILL_METHOD] + ) + self.prompt_method = self.methods.get(PREFILL_METHOD) or self.step_method + first = programs[0] + self.eos_ids = constant_method(first, "get_eos_ids", []) + if not isinstance(self.eos_ids, list): + self.eos_ids = [self.eos_ids] + self.logits_scale = constant_method(first, "get_logits_scale", 1.0) + self.logits_zero_point = constant_method(first, "get_logits_zero_point", 0) + self.num_layers = self.step_method.num_layers + self.max_seq_len = self.step_method.context_len + # Past caches by past length: methods with different step sizes lay + # the same positions out with different strides. + self.caches: Dict[int, List[torch.Tensor]] = {} + for method in self.methods.values(): + if method.past_len not in self.caches: + self.caches[method.past_len] = [ + torch.zeros(shape, dtype=dtype) + for shape, dtype in method.cache_shapes + ] + self.filled = 0 + self.active_past_len: Optional[int] = None + + def reset(self) -> None: + for cache in self.caches.values(): + for tensor in cache: + tensor.zero_() + self.filled = 0 + self.active_past_len = None + + def past(self, past_len: Optional[int] = None) -> List[torch.Tensor]: + """The K caches then V caches in the layout of ``past_len`` (default: + the one holding the filled positions).""" + return self.caches[self.active_past_len if past_len is None else past_len] + + def _rearrange(self, past_len: int) -> None: + src = self.active_past_len + if src is not None and src != past_len and self.filled: + self.adopt(self.caches[src], self.filled, past_len) + self.active_past_len = past_len + + def adopt( + self, + caches: Sequence[torch.Tensor], + filled: int, + past_len: Optional[int] = None, + ) -> None: + """Copy ``filled`` positions of caches in another layout into the + layout of ``past_len`` (default: the single-token step's).""" + past_len = self.step_method.past_len if past_len is None else past_len + dst = self.caches[past_len] + n = self.num_layers + for i in range(n): + dst[i][..., :filled] = caches[i][..., :filled] + dst[n + i][..., :filled, :] = caches[n + i][..., :filled, :] + self.filled = filled + self.active_past_len = past_len + + def _step(self, method: _Method, chunk: List[int]) -> torch.Tensor: + n, start = len(chunk), self.filled + if start + n > method.past_len + (1 if method.ar_len == 1 else 0): + raise ValueError( + f"Context full: position {start} with {method.past_len} slots of past." + ) + padded = chunk + [chunk[-1]] * (method.ar_len - n) + method.tokens.view(-1)[:] = torch.tensor(padded, dtype=method.tokens.dtype) + method.positions.view(-1)[:] = torch.tensor( + [min(start + i, method.context_len - 1) for i in range(method.ar_len)], + dtype=method.positions.dtype, + ) + self._rearrange(method.past_len) + method.fill_mask(start, n) + caches = self.caches[method.past_len] + outputs = method.method.execute( + [method.tokens, method.mask, method.positions, *caches] + ) + layers = self.num_layers + for i in range(layers): + caches[i][..., start : start + n] = outputs[1 + i][..., :n] + caches[layers + i][..., start : start + n, :] = outputs[1 + layers + i][ + ..., :n, : + ] + self.filled = start + n + logits = outputs[0].reshape(-1, outputs[0].shape[-1])[n - 1] + if not logits.dtype.is_floating_point: + logits = ( + logits.to(torch.int64) - self.logits_zero_point + ).float() * self.logits_scale + return logits + + def forward(self, tokens: Sequence[int]) -> torch.Tensor: + """Run ``tokens`` after the filled positions; the last token's logits.""" + tokens = list(tokens) + if not tokens: + raise ValueError("No tokens to run.") + method = self.prompt_method if len(tokens) > 1 else self.step_method + logits = None + for i in range(0, len(tokens), method.ar_len): + logits = self._step(method, tokens[i : i + method.ar_len]) + return logits + + def generate(self, prompt_ids: Sequence[int], max_new_tokens: int) -> List[int]: + """Greedy continuation of ``prompt_ids`` (new tokens only), stopping + at an end-of-sequence id or a full context.""" + self.reset() + logits = self.forward(prompt_ids) + generated: List[int] = [] + while len(generated) < max_new_tokens: + token = int(torch.argmax(logits)) + generated.append(token) + if token in self.eos_ids or self.filled >= self.max_seq_len: + break + logits = self.forward([token]) + return generated diff --git a/extension/llm/io_contracts/test/__init__.py b/extension/llm/io_contracts/test/__init__.py new file mode 100644 index 00000000000..d35f6369753 --- /dev/null +++ b/extension/llm/io_contracts/test/__init__.py @@ -0,0 +1,7 @@ +# 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. + +"""I/O contract regression test package.""" diff --git a/extension/llm/io_contracts/test/test_hf_static.py b/extension/llm/io_contracts/test/test_hf_static.py new file mode 100644 index 00000000000..958e7f1a7c7 --- /dev/null +++ b/extension/llm/io_contracts/test/test_hf_static.py @@ -0,0 +1,503 @@ +# 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. + +"""Exercise the registered hf_static runner on real tiny HF and ET programs.""" + +import unittest +from types import SimpleNamespace +from unittest.mock import patch + +import torch +from executorch.exir import EdgeCompileConfig, ExecutorchBackendConfig, to_edge +from executorch.exir.passes import MemoryPlanningPass +from executorch.extension.llm.export import hf_static +from executorch.extension.llm.io_contracts.hf_static import ( + HfStaticDriver, + HfStaticRunner, +) +from executorch.extension.pybindings.portable_lib import ( + _load_for_executorch_from_buffer, +) + +try: + from transformers import GenerationConfig, MistralConfig, MistralForCausalLM + from transformers.exporters import ( + DynamoConfig, + DynamoExporter, + ExecutorchModelRunner, + ExportedGenerator, + ) + from transformers.exporters.io_contracts import io_contract_runner +except ImportError: + DynamoExporter = None + + +class _PortableProgram: + def __init__(self, buffer): + self.module = _load_for_executorch_from_buffer(buffer) + self.method_names = set(self.module.method_names()) + + def load_method(self, name): + return SimpleNamespace( + metadata=self.module.method_meta(name), + execute=lambda inputs: self.module.run_method(name, inputs), + ) + + +class _Program: + def __init__(self, program, constants=None, methods=None): + self.program = program + self.constants = constants or {} + self.methods = methods or {} + self.method_names = program.method_names + + def load_method(self, name): + if name in self.constants: + return SimpleNamespace(execute=lambda inputs: [self.constants[name]]) + return self.methods.get(name) or self.program.load_method(name) + + +@unittest.skipIf(DynamoExporter is None, "needs transformers.exporters") +class TestHfStaticRunner(unittest.TestCase): + CONTEXT = 32 + WINDOW = 6 + CHUNK = 4 + PROMPT = [5, 17, 42, 9, 77, 3, 64, 33, 21, 8, 99] + + @classmethod + def setUpClass(cls): + torch.manual_seed(0) + cls.model = MistralForCausalLM( + MistralConfig( + hidden_size=64, + intermediate_size=128, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + vocab_size=128, + max_position_embeddings=cls.CONTEXT, + sliding_window=cls.WINDOW, + eos_token_id=None, + ) + ).eval() + sample = torch.tensor([[3, 4, 5, 6]]) + cls.generation_config = GenerationConfig( + do_sample=False, + cache_implementation="static", + max_cache_len=cls.CONTEXT, + ) + with torch.no_grad(): + artifacts = DynamoExporter().export_for_generation( + cls.model, + {"input_ids": sample, "attention_mask": torch.ones_like(sample)}, + DynamoConfig(dynamic=True), + generation_config=cls.generation_config, + keep_all_logits=True, + ) + cls.hf_graph = artifacts.components["decode"].artifact + cls.metadata = artifacts.components["decode"].metadata.raw + cls.programs = { + name: hf_static.to_hf_static( + hf_static.specialize( + cls.hf_graph, q, cls.CONTEXT, config=cls.model.config + ) + ) + for name, q in (("decode", 1), ("prefill", cls.CHUNK)) + } + constants = { + "get_io_contract": "hf_static", + "get_io_contract_version": 1, + "get_max_seq_len": cls.CONTEXT, + "get_n_layers": 2, + "get_cache_dtype": "float32", + "get_eos_ids": [], + } + cls.et = to_edge( + cls.programs, + compile_config=EdgeCompileConfig(_check_ir_validity=False), + constant_methods=constants, + ).to_executorch( + ExecutorchBackendConfig( + run_reinplace_pass=True, + memory_planning_pass=MemoryPlanningPass( + alloc_graph_input=False, alloc_graph_output=False + ), + ) + ) + cls.buffer = cls.et.buffer + + def _load(self): + return _PortableProgram(self.buffer) + + def _expected(self, prompt, count): + with torch.no_grad(): + return self.model.generate( + torch.tensor([prompt]), max_new_tokens=count, do_sample=False + )[0, len(prompt) :].tolist() + + def test_all_sliding_mistral_shapes_and_greedy(self): + layout = hf_static.HfStaticLayout.of(self.programs["decode"]) + self.assertEqual(layout.max_seq_len, self.CONTEXT) + self.assertEqual([s[2] for s in layout.cache_shapes], [self.WINDOW] * 4) + runner = HfStaticRunner(self._load()) + self.assertEqual( + runner.generate(self.PROMPT, 20), self._expected(self.PROMPT, 20) + ) + + def test_native_hf_loader_and_exported_generator(self): + runner = ExecutorchModelRunner._load(self.buffer, self.metadata) + generator = ExportedGenerator( + {"decoder": runner}, self.model.config, self.generation_config + ) + for prompt in (self.PROMPT, [11]): + with self.subTest(prompt=prompt): + output = generator.generate( + input_ids=torch.tensor([prompt]), max_new_tokens=8, do_sample=False + ) + self.assertEqual( + output[0, len(prompt) :].tolist(), self._expected(prompt, 8) + ) + + def test_exported_generator_uses_registered_runner(self): + runner = io_contract_runner( + "hf_static", "executorch", self._load(), export_metadata=self.metadata + ) + self.assertTrue(runner.drives_generation) + generator = ExportedGenerator( + {"decoder": runner}, self.model.config, self.generation_config + ) + for prompt in (self.PROMPT, [11]): + with self.subTest(prompt=prompt): + result = generator.generate( + input_ids=torch.tensor([prompt]), max_new_tokens=8, do_sample=False + ) + self.assertEqual( + result[0, len(prompt) :].tolist(), self._expected(prompt, 8) + ) + + def test_exported_generator_rejects_caller_padding_and_custom_positions(self): + runner = io_contract_runner( + "hf_static", "executorch", self._load(), export_metadata=self.metadata + ) + generator = ExportedGenerator( + {"decoder": runner}, self.model.config, self.generation_config + ) + ids = torch.tensor([[0, *self.PROMPT]]) + mask = torch.tensor([[0, *([1] * len(self.PROMPT))]]) + with self.assertRaisesRegex(ValueError, "caller-padded"): + generator.generate( + input_ids=ids, attention_mask=mask, max_new_tokens=2, do_sample=False + ) + runner.reset_state() + with self.assertRaisesRegex(ValueError, "consecutive absolute"): + runner( + input_ids=torch.tensor([self.PROMPT]), + position_ids=torch.arange(len(self.PROMPT))[None] + 1, + ) + + def test_padding_rows_are_false_when_window_equals_context(self): + captured = [] + layout = hf_static.HfStaticLayout( + [ + "input_ids", + "attention_mask", + "position_ids", + "past_key_values_0", + "past_key_values_1", + "past_key_values_2", + ], + 4, + 32, + 1, + [torch.Size([1, 1, 32, 2])] * 2, + [torch.Size([1, 1, 32, 2]), torch.Size([1, 1, 32, 2]), torch.Size([])], + torch.float32, + {"attention_mask": 32}, + state_dtypes=[torch.float32, torch.float32, torch.int64], + ) + + def step(*inputs): + captured.append(inputs[1]) + return [torch.zeros(1, 4, 2), *inputs[-3:]] + + driver = HfStaticDriver({4: step}, layout) + driver.forward([1, 2, 3], 0, 4) + self.assertFalse(bool(captured[0][:, :, 3].any())) + self.assertTrue(bool(captured[0][:, :, :3].any())) + + def test_registered_runner_keeps_cache_buffers_across_steps(self): + runner = HfStaticRunner(self._load()) + pointers = [tensor.data_ptr() for tensor in runner.driver.state] + runner.forward(self.PROMPT) + self.assertEqual( + pointers, [tensor.data_ptr() for tensor in runner.driver.state] + ) + runner.forward([121]) + self.assertEqual( + pointers, [tensor.data_ptr() for tensor in runner.driver.state] + ) + + def test_reset_and_padded_chunk_counters(self): + runner = HfStaticRunner(self._load()) + first = runner.generate(self.PROMPT, 8) + self.assertEqual( + [t.item() for t in runner.driver.state if t.dim() == 0], + [len(self.PROMPT) + 8] * 2, + ) + pointers = [t.data_ptr() for t in runner.driver.state] + runner.reset() + self.assertEqual( + [t.item() for t in runner.driver.state if t.dim() == 0], [0, 0] + ) + self.assertEqual(pointers, [t.data_ptr() for t in runner.driver.state]) + self.assertEqual(runner.generate(self.PROMPT, 8), first) + + def test_prepare_overwrites_stale_counters(self): + runner = HfStaticRunner(self._load()) + for tensor in runner.driver.state: + if tensor.dim() == 0: + tensor.fill_(999) + runner.forward(self.PROMPT) + self.assertEqual( + [t.item() for t in runner.driver.state if t.dim() == 0], + [len(self.PROMPT)] * 2, + ) + + def test_near_context_tail_uses_decode(self): + runner = HfStaticRunner(self._load()) + prefix = (self.PROMPT * 3)[:29] + runner.forward(prefix) + calls = [] + for q, step in list(runner.driver.steps.items()): + + def record(*inputs, q=q, step=step): + calls.append(q) + return step(*inputs) + + runner.driver.steps[q] = record + logits = runner.forward([3, 4]) + self.assertEqual(calls, [1, 1]) + self.assertEqual(runner.filled, 31) + expected = self._expected(prefix + [3, 4], 1)[0] + self.assertEqual(int(logits.argmax()), expected) + runner.forward([expected]) + with self.assertRaisesRegex(ValueError, "remaining context"): + runner.forward([1]) + + def test_cache_aliases_and_no_functional_cache_copy(self): + for plan in self.et.executorch_program.execution_plan: + if plan.name not in ("prefill", "decode"): + continue + self.assertEqual(sorted(plan.outputs[:6]), sorted(plan.inputs[-6:])) + ops = [o.name for o in plan.operators] + self.assertNotIn("aten::index_put", ops) + self.assertNotIn("aten::index_copy", ops) + instructions = [ + i.instr_args + for c in plan.chains + for i in c.instructions + if hasattr(i.instr_args, "op_index") + ] + writes = [ + i + for i in instructions + if ops[i.op_index] in ("aten::index_put_", "aten::index_copy_") + ] + self.assertEqual(len(writes), 8) + for write in writes: + self.assertIn(write.args[0], plan.inputs) + copies = [i for i in instructions if ops[i.op_index] == "aten::copy_"] + self.assertEqual(len(copies), 2) + for copy in copies: + self.assertEqual(plan.values[copy.args[0]].val.sizes, []) + layout = hf_static.HfStaticLayout.of(self.programs["decode"]) + module = _load_for_executorch_from_buffer(self.buffer) + steps = {} + for q, method_name in ((1, "decode"), (4, "prefill")): + + def step(*inputs, method_name=method_name): + outputs = module.run_method(method_name, inputs, clone_outputs=False) + self.assertEqual(len(outputs), 13) + self.assertEqual( + sorted(t.data_ptr() for t in inputs[-6:]), + sorted(t.data_ptr() for t in outputs[:6]), + ) + return [outputs[6], *outputs[7:]] + + steps[q] = step + driver = HfStaticDriver(steps, layout, in_place=True) + self.assertEqual( + driver.generate(self.PROMPT, 8)[len(self.PROMPT) :], + self._expected(self.PROMPT, 8), + ) + + def test_rejects_newer_or_invalid_versions(self): + for version in (2, 0, None, True): + with self.subTest(version=version), self.assertRaisesRegex( + ValueError, "contract version" + ): + HfStaticRunner( + _Program(self._load(), {"get_io_contract_version": version}) + ) + + def test_rejects_wrong_cache_dtype(self): + for dtype in ("float16", "uint16", None): + with self.subTest(dtype=dtype), self.assertRaisesRegex( + ValueError, "get_cache_dtype" + ): + HfStaticRunner(_Program(self._load(), {"get_cache_dtype": dtype})) + + def test_rejects_mismatched_layer_or_context_metadata(self): + for constants in ( + {"get_n_layers": 3}, + {"get_max_seq_len": 4}, + {"get_n_layers": 0}, + ): + with self.subTest(constants=constants), self.assertRaises(ValueError): + HfStaticRunner(_Program(self._load(), constants)) + + def test_cpu_tags_survive_export_and_to_edge(self): + edge = to_edge( + self.programs, compile_config=EdgeCompileConfig(_check_ir_validity=False) + ) + for name in self.programs: + hf_static.tag_cache_mutations(edge.exported_program(name)) + graph = edge.exported_program(name).graph + writes = [n for n in graph.nodes if "index_put" in str(n.target)] + self.assertEqual(len(writes), 8) + self.assertTrue( + all(n.meta.get("custom", {}).get(hf_static.CPU_META) for n in writes) + ) + booleans = [ + n + for n in graph.nodes + if n.op == "call_function" + and isinstance(n.meta.get("val"), torch.Tensor) + and n.meta["val"].dtype == torch.bool + ] + self.assertTrue(booleans) + self.assertTrue( + all(n.meta.get("custom", {}).get(hf_static.CPU_META) for n in booleans) + ) + + def test_partial_attention_match_is_rejected(self): + program = hf_static.specialize( + self.hf_graph, 1, self.CONTEXT, config=self.model.config + ) + original = hf_static._cache_attention + with patch.object( + hf_static, "_cache_attention", side_effect=lambda graph: original(graph)[:1] + ): + with self.assertRaisesRegex(ValueError, "Expected 2.*matched 1"): + hf_static.to_hf_static(program) + + def test_rejects_malformed_prefill_inputs(self): + for index, shape, dtype in ( + (0, (1, 4), torch.int32), + (1, (1, 1, 1, 6), torch.bool), + (2, (1, 1), torch.int64), + (2, (1, 4), torch.int32), + (3, (1, 1, 6, 32), torch.float32), + (5, (), torch.int32), + ): + program = self._load() + method = program.load_method("prefill") + original = method.metadata + + class Metadata: + def __init__(self, original, index, shape, dtype): + self.original, self.index = original, index + self.shape, self.dtype = shape, dtype + + def __getattr__(self, name): + return getattr(self.original, name) + + def input_tensor_meta(self, i): + return ( + SimpleNamespace( + sizes=lambda: self.shape, dtype=lambda: self.dtype + ) + if i == self.index + else self.original.input_tensor_meta(i) + ) + + bad_method = SimpleNamespace( + metadata=Metadata(original, index, shape, dtype) + ) + with self.subTest(index=index), self.assertRaises(ValueError): + HfStaticRunner(_Program(program, methods={"prefill": bad_method})) + + def test_rejects_int32_tokens_or_positions_when_methods_agree(self): + program = self._load() + for index in (0, 2): + + class Metadata: + def __init__(self, original, index): + self.original, self.index = original, index + + def __getattr__(self, name): + return getattr(self.original, name) + + def input_tensor_meta(self, i): + info = self.original.input_tensor_meta(i) + return ( + SimpleNamespace(sizes=info.sizes, dtype=lambda: torch.int32) + if i == self.index + else info + ) + + methods = { + name: SimpleNamespace( + metadata=Metadata(program.load_method(name).metadata, index) + ) + for name in ("decode", "prefill") + } + with self.subTest(index=index), self.assertRaisesRegex( + ValueError, "input_ids|position_ids" + ): + HfStaticRunner(_Program(program, methods=methods)) + + def test_rejects_float64_even_with_matching_cache_metadata(self): + program = self._load() + + class Metadata: + def __init__(self, original): + self.original = original + + def __getattr__(self, name): + return getattr(self.original, name) + + def input_tensor_meta(self, i): + info = self.original.input_tensor_meta(i) + return ( + SimpleNamespace(sizes=info.sizes, dtype=lambda: torch.float64) + if i in (3, 4, 6, 7) + else info + ) + + def output_tensor_meta(self, i): + info = self.original.output_tensor_meta(i) + return ( + SimpleNamespace(sizes=info.sizes, dtype=lambda: torch.float64) + if len(info.sizes()) == 4 + else info + ) + + methods = { + name: SimpleNamespace(metadata=Metadata(program.load_method(name).metadata)) + for name in ("decode", "prefill") + } + with self.assertRaisesRegex(ValueError, "float cache inputs"): + HfStaticRunner(_Program(program, {"get_cache_dtype": "float64"}, methods)) + + def test_zero_new_tokens(self): + runner = HfStaticRunner(self._load()) + self.assertEqual(runner.generate(self.PROMPT, 0), []) + self.assertEqual(runner.filled, len(self.PROMPT)) + + +if __name__ == "__main__": + unittest.main() diff --git a/extension/llm/io_contracts/test/test_kv_cache_io.py b/extension/llm/io_contracts/test/test_kv_cache_io.py new file mode 100644 index 00000000000..355c9981770 --- /dev/null +++ b/extension/llm/io_contracts/test/test_kv_cache_io.py @@ -0,0 +1,109 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# All rights reserved. +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +"""The kvcache_io runner's own protocol, without a QNN export dependency.""" + +import unittest + +import torch +from executorch.exir import to_edge_transform_and_lower +from executorch.extension.llm.io_contracts import io_contract, KvCacheIORunner +from executorch.runtime import Runtime, Verification + +try: + from transformers import GenerationConfig, LlamaConfig + from transformers.exporters import ExecutorchModelRunner, ExportedGenerator +except ImportError: + ExportedGenerator = None + +L = 32 +PROMPTS = ([5, 17, 42, 9, 77, 3, 64], [11]) + + +class _Step(torch.nn.Module): + def forward(self, tokens, mask, positions, past_k, past_v): + rows = tokens.float().view(1, 1, -1, 1) + keys = torch.cat((past_k.transpose(2, 3), rows), 2) + values = torch.cat((past_v, rows), 2) + history = (keys + values).squeeze(-1).unsqueeze(2) + next_id = (history * (mask == 0)).sum(-1).squeeze(1) % 128 + logits = -(torch.arange(128).float()[None, None] - next_id[..., None]).abs() + return logits, rows.transpose(2, 3), rows + + +class TestKvCacheIORunner(unittest.TestCase): + @classmethod + def setUpClass(cls): + steps = {} + for query in (1, 4): + inputs = ( + torch.zeros(1, query, dtype=torch.int32), + torch.zeros(1, 1, query, L), + torch.zeros(1, query, dtype=torch.int32), + torch.zeros(1, 1, 1, L - query), + torch.zeros(1, 1, L - query, 1), + ) + steps[query] = torch.export.export(_Step(), inputs) + cls.programs = {} + for name, methods in ( + ("decode", {"kv_forward": steps[1]}), + ("hybrid", {"kv_forward": steps[1], "prefill_forward": steps[4]}), + ): + cls.programs[name] = ( + to_edge_transform_and_lower( + methods, + constant_methods={ + "get_io_contract": "kvcache_io", + "get_io_contract_version": 1, + "get_max_seq_len": L, + "get_n_layers": 1, + }, + ) + .to_executorch() + .buffer + ) + + def _load(self, name): + return Runtime.get().load_program( + self.programs[name], verification=Verification.Minimal + ) + + def _expected(self, prompt, new_tokens): + sequence = list(prompt) + for _ in range(new_tokens): + sequence.append(2 * sum(sequence) % 128) + return sequence[len(prompt) :] + + def test_greedy_matches_model(self): + for name in ("decode", "hybrid"): + program = self._load(name) + self.assertEqual(io_contract(program), "kvcache_io") + runner = KvCacheIORunner([program]) + for prompt in PROMPTS: + with self.subTest(program=name, prompt=prompt): + self.assertEqual( + runner.generate(prompt, 8), self._expected(prompt, 8) + ) + + @unittest.skipIf(ExportedGenerator is None, "needs transformers.exporters") + def test_exported_generator_drives_the_contract(self): + runner = ExecutorchModelRunner._load(self.programs["hybrid"], {}) + config = LlamaConfig(vocab_size=128, eos_token_id=None) + generator = ExportedGenerator( + {"decoder": runner}, config, GenerationConfig(do_sample=False) + ) + for prompt in PROMPTS: + with self.subTest(prompt=prompt): + out = generator.generate( + input_ids=torch.tensor([prompt]), max_new_tokens=8, do_sample=False + ) + self.assertEqual( + out[0, len(prompt) :].tolist(), self._expected(prompt, 8) + ) + + +if __name__ == "__main__": + unittest.main() diff --git a/extension/llm/runner/constants.h b/extension/llm/runner/constants.h index 60a6383a72e..84aecc71bd6 100644 --- a/extension/llm/runner/constants.h +++ b/extension/llm/runner/constants.h @@ -14,6 +14,9 @@ inline constexpr auto kEnableDynamicShape = "enable_dynamic_shape"; inline constexpr auto kBosId = "get_bos_id"; inline constexpr auto kEosIds = "get_eos_ids"; inline constexpr auto kMaxSeqLen = "get_max_seq_len"; +// Names the I/O contract of the program's methods (e.g. "kvcache_io"), +// which selects the IOManager that drives them. +inline constexpr auto kIoContract = "get_io_contract"; inline constexpr auto kMaxContextLen = "get_max_context_len"; inline constexpr auto kVocabSize = "get_vocab_size"; inline constexpr auto kActivationDtype = "get_activation_dtype"; diff --git a/extension/llm/runner/io_manager/hf_static_io_manager.cpp b/extension/llm/runner/io_manager/hf_static_io_manager.cpp new file mode 100644 index 00000000000..081ce85245e --- /dev/null +++ b/extension/llm/runner/io_manager/hf_static_io_manager.cpp @@ -0,0 +1,538 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include + +#include + +#include +#include +#include +#include +#include +#include + +namespace executorch { +namespace extension { +namespace llm { + +using ::executorch::aten::ScalarType; +using ::executorch::runtime::Error; +using ::executorch::runtime::EValue; +using ::executorch::runtime::Result; +using ::executorch::runtime::TensorInfo; + +namespace { +constexpr const char* kForward = "forward"; +constexpr size_t kAlignment = 64; + +size_t align_up(size_t n) { + return (n + kAlignment - 1) / kAlignment * kAlignment; +} + +} // namespace + +HfStaticIOManager::HfStaticIOManager(ET_MODULE_NAMESPACE::Module& module) + : IOManager(module), module_(module) { + if (auto& factory = io_buffer_allocator_factory()) { + allocator_ = factory(); + } +} + +HfStaticIOManager::~HfStaticIOManager() { + methods_.clear(); + caches_.clear(); + for (const auto& block : blocks_) { + if (block.from_allocator) { + allocator_->deallocate(block.data); + } else { + ::operator delete(block.data, std::align_val_t(kAlignment)); + } + } +} + +std::pair HfStaticIOManager::method_names( + const std::unordered_set& available) const { + if (available.count(kDecodeMethod)) { + return std::make_pair( + available.count(kPrefillMethod) ? kPrefillMethod : kDecodeMethod, + kDecodeMethod); + } + return available.count(kForward) ? std::make_pair(kForward, kForward) + : std::make_pair("", ""); +} + +int64_t HfStaticIOManager::max_tokens_per_step( + const std::string& method) const { + auto it = methods_.find(method); + return it == methods_.end() ? 1 : it->second.query_len; +} + +int64_t HfStaticIOManager::max_tokens_per_step( + const std::string& method, + int64_t start_pos) const { + auto it = methods_.find(method); + if (it == methods_.end() || start_pos < 0 || + start_pos > context_len_ - it->second.query_len) { + return 0; + } + return it->second.query_len; +} + +void HfStaticIOManager::allocate( + const std::vector>& buffers) { + // One block for all of them: a delegate registers shared memory per block. + size_t total = 0; + for (const auto& [buf, info] : buffers) { + total = align_up(total) + info.nbytes(); + } + void* block = allocator_ ? allocator_->allocate(total, kAlignment) : nullptr; + const bool from_allocator = block != nullptr; + if (!from_allocator) { + block = ::operator new(total, std::align_val_t(kAlignment)); + } + blocks_.push_back({block, from_allocator}); + std::memset(block, 0, total); + size_t offset = 0; + for (const auto& [buf, info] : buffers) { + offset = align_up(offset); + buf->data = static_cast(block) + offset; + buf->nbytes = info.nbytes(); + std::vector sizes( + info.sizes().begin(), info.sizes().end()); + buf->tensor = from_blob(buf->data, sizes, info.scalar_type()); + if (from_allocator) { + allocator_->bind(block, buf->data); + } + offset += info.nbytes(); + } +} + +Error HfStaticIOManager::load_metadata() { + if (metadata_loaded_) { + return Error::Ok; + } + auto names = module_.method_names(); + ET_CHECK_OK_OR_RETURN_ERROR(names.error()); + for (const auto* name : + {"get_io_contract_version", + "get_cache_dtype", + "get_n_layers", + "get_max_seq_len"}) { + ET_CHECK_OR_RETURN_ERROR( + names->count(name), + InvalidArgument, + "hf_static: required metadata method %s is missing.", + name); + } + auto version = module_.get("get_io_contract_version"); + ET_CHECK_OK_OR_RETURN_ERROR(version.error()); + ET_CHECK_OR_RETURN_ERROR( + version->isInt() && version->toInt() >= 1, + InvalidArgument, + "hf_static: get_io_contract_version must be a positive integer."); + ET_CHECK_OR_RETURN_ERROR( + version->toInt() <= kContractVersion, + NotSupported, + "hf_static: contract version %" PRId64 + " is newer than supported version %" PRId64 ".", + version->toInt(), + kContractVersion); + auto dtype = module_.get("get_cache_dtype"); + ET_CHECK_OK_OR_RETURN_ERROR(dtype.error()); + ET_CHECK_OR_RETURN_ERROR( + dtype->isString(), + InvalidArgument, + "hf_static: get_cache_dtype must be float32 or float16."); + const std::string label(dtype->toString()); + ET_CHECK_OR_RETURN_ERROR( + label == "float32" || label == "float16", + InvalidArgument, + "hf_static: unsupported cache dtype '%s'.", + label.c_str()); + auto layers = module_.get("get_n_layers"); + ET_CHECK_OK_OR_RETURN_ERROR(layers.error()); + auto context = module_.get("get_max_seq_len"); + ET_CHECK_OK_OR_RETURN_ERROR(context.error()); + ET_CHECK_OR_RETURN_ERROR( + layers->isInt() && layers->toInt() > 0 && + layers->toInt() <= std::numeric_limits::max() && + context->isInt() && context->toInt() > 0 && + context->toInt() <= std::numeric_limits::max(), + InvalidArgument, + "hf_static: get_n_layers and get_max_seq_len must be positive integers."); + cache_dtype_ = label == "float32" ? ScalarType::Float : ScalarType::Half; + num_layers_ = static_cast(layers->toInt()); + context_len_ = context->toInt(); + metadata_loaded_ = true; + return Error::Ok; +} + +Error HfStaticIOManager::load_method(const std::string& method) { + ET_CHECK_OK_OR_RETURN_ERROR(load_metadata()); + if (methods_.count(method)) { + return Error::Ok; + } + ET_CHECK_OK_OR_RETURN_ERROR(module_.load_method(method)); + auto meta_result = module_.method_meta(method); + ET_CHECK_OK_OR_RETURN_ERROR(meta_result.error()); + auto meta = std::move(*meta_result); + Method io; + io.num_inputs = meta.num_inputs(); + std::vector infos; + for (size_t i = 0; i < io.num_inputs; ++i) { + auto info = meta.input_tensor_meta(i); + ET_CHECK_OK_OR_RETURN_ERROR(info.error()); + infos.push_back(*info); + } + auto is = [&](size_t i, ScalarType dtype, size_t rank) { + return i < infos.size() && infos[i].scalar_type() == dtype && + infos[i].sizes().size() == rank; + }; + // input_ids, the attention mask(s), position_ids, then per layer keys, + // values and the fill counter. + size_t i = 0; + ET_CHECK_OR_RETURN_ERROR( + is(i, ScalarType::Long, 2), + InvalidArgument, + "hf_static: %s: input 0 must be int64 input_ids [1, N].", + method.c_str()); + io.input_ids = i++; + const auto token_sizes = infos[io.input_ids].sizes(); + ET_CHECK_OR_RETURN_ERROR( + token_sizes[0] == 1 && token_sizes[1] > 0, + InvalidArgument, + "hf_static: %s: input_ids must be [1, N] with N > 0.", + method.c_str()); + io.query_len = token_sizes[1]; + while (is(i, ScalarType::Bool, 4)) { + const auto sizes = infos[i].sizes(); + ET_CHECK_OR_RETURN_ERROR( + sizes[0] == 1 && sizes[1] == 1 && sizes[2] == io.query_len && + sizes[3] > 0 && sizes[3] <= context_len_, + InvalidArgument, + "hf_static: %s: mask %zu must be [1, 1, N, W], 0 < W <= context.", + method.c_str(), + i); + io.masks.push_back(i++); + } + ET_CHECK_OR_RETURN_ERROR( + !io.masks.empty() && is(i, ScalarType::Long, 2), + InvalidArgument, + "hf_static: %s: expected bool attention masks, then int64 position_ids.", + method.c_str()); + io.position_ids = i++; + const auto position_sizes = infos[io.position_ids].sizes(); + ET_CHECK_OR_RETURN_ERROR( + position_sizes[0] == 1 && position_sizes[1] == io.query_len, + InvalidArgument, + "hf_static: %s: position_ids must be [1, N].", + method.c_str()); + io.first_cache = i; + const size_t num_cache_inputs = io.num_inputs - io.first_cache; + ET_CHECK_OR_RETURN_ERROR( + num_cache_inputs == 3 * static_cast(num_layers_), + InvalidArgument, + "hf_static: %s: expected three cache inputs per layer, matching get_n_layers.", + method.c_str()); + for (size_t c = io.first_cache; c < io.num_inputs; c += 3) { + ET_CHECK_OR_RETURN_ERROR( + is(c, cache_dtype_, 4) && is(c + 1, cache_dtype_, 4) && + is(c + 2, ScalarType::Long, 0), + InvalidArgument, + "hf_static: %s: inputs %zu..%zu must be float keys/values matching get_cache_dtype and an int64 counter.", + method.c_str(), + c, + c + 2); + const auto keys = infos[c].sizes(); + const auto values = infos[c + 1].sizes(); + ET_CHECK_OR_RETURN_ERROR( + keys[0] == 1 && keys[1] > 0 && keys[2] >= io.query_len && + keys[2] <= context_len_ && keys[3] > 0 && + std::equal(keys.begin(), keys.end(), values.begin()), + InvalidArgument, + "hf_static: %s: keys and values must share [1, H, L, D], N <= L <= context.", + method.c_str()); + } + for (size_t s = 0; s < io.first_cache; ++s) { + io.step_inputs[s]; // Allocated once the method is stored, below. + } + ET_CHECK_OR_RETURN_ERROR( + meta.num_outputs() == num_cache_inputs + 1 || + meta.num_outputs() == 2 * num_cache_inputs + 1, + InvalidArgument, + "hf_static: %s must return mutated inputs, logits, and optionally the user cache state.", + method.c_str()); + if (meta.num_outputs() == 2 * num_cache_inputs + 1) { + for (size_t c = 0; c < num_cache_inputs; ++c) { + auto state = meta.output_tensor_meta(num_cache_inputs + 1 + c); + ET_CHECK_OK_OR_RETURN_ERROR(state.error()); + const auto& input = infos[io.first_cache + c]; + ET_CHECK_OR_RETURN_ERROR( + state->scalar_type() == input.scalar_type() && + state->sizes().size() == input.sizes().size() && + std::equal( + input.sizes().begin(), + input.sizes().end(), + state->sizes().begin()), + InvalidArgument, + "hf_static: %s: user state output %zu differs from its input.", + method.c_str(), + c); + } + } + const auto logits_meta = meta.output_tensor_meta(num_cache_inputs); + ET_CHECK_OK_OR_RETURN_ERROR(logits_meta.error()); + const auto logits_sizes = logits_meta->sizes(); + ET_CHECK_OR_RETURN_ERROR( + (logits_meta->scalar_type() == ScalarType::Float || + logits_meta->scalar_type() == ScalarType::Half) && + logits_sizes.size() == 3 && logits_sizes[0] == 1 && + logits_sizes[1] == io.query_len && logits_sizes[2] > 0, + InvalidArgument, + "hf_static: %s: logits must be float32/float16 [1, N, V].", + method.c_str()); + if (caches_.empty()) { + caches_.resize(num_cache_inputs); + std::vector> cache_buffers; + for (size_t c = 0; c < num_cache_inputs; ++c) { + cache_buffers.emplace_back(&caches_[c], infos[io.first_cache + c]); + } + allocate(cache_buffers); + logits_f_.assign(logits_sizes[2], 0.0f); + logits_f_tensor_ = from_blob( + logits_f_.data(), + {1, 1, static_cast(logits_f_.size())}, + ScalarType::Float); + } else { + ET_CHECK_OR_RETURN_ERROR( + num_cache_inputs == caches_.size() && + static_cast(logits_sizes[2]) == logits_f_.size(), + InvalidArgument, + "hf_static: %s takes a different cache or vocabulary than the other method.", + method.c_str()); + for (size_t c = 0; c < caches_.size(); ++c) { + const auto& info = infos[io.first_cache + c]; + const auto& tensor = *caches_[c].tensor; + ET_CHECK_OR_RETURN_ERROR( + info.scalar_type() == tensor.scalar_type() && + info.sizes().size() == tensor.sizes().size() && + std::equal( + info.sizes().begin(), + info.sizes().end(), + tensor.sizes().begin()), + InvalidArgument, + "hf_static: %s: cache input %zu differs from the other method's dtype or shape.", + method.c_str(), + c); + } + } + // Every other output is an input written in place. Unplanned logits are + // written into a buffer of ours (shared memory, so a delegate writes them + // directly); planned ones come back in the program's own memory. + const bool bind_logits = !logits_meta->is_memory_planned(); + auto& stored = methods_.emplace(method, std::move(io)).first->second; + std::vector> method_buffers; + for (auto& [s, buf] : stored.step_inputs) { + method_buffers.emplace_back(&buf, infos[s]); + } + if (bind_logits) { + method_buffers.emplace_back(&stored.logits, *logits_meta); + } + allocate(method_buffers); + if (meta.num_outputs() == 2 * num_cache_inputs + 1) { + for (size_t c = 0; c < num_cache_inputs; ++c) { + const size_t output = num_cache_inputs + 1 + c; + if (!meta.output_tensor_meta(output)->is_memory_planned()) { + ET_CHECK_OK_OR_RETURN_ERROR( + module_.set_output(method, EValue(*caches_[c].tensor), output)); + } + } + } + if (bind_logits) { + ET_CHECK_OK_OR_RETURN_ERROR(module_.set_output( + method, EValue(*stored.logits.tensor), num_cache_inputs)); + } + return Error::Ok; +} + +Error HfStaticIOManager::load( + const std::string& prefill_method, + const std::string& decode_method) { + ET_CHECK_OK_OR_RETURN_ERROR(load_method(prefill_method)); + return load_method(decode_method); +} + +Error HfStaticIOManager::reset( + const std::string& prefill_method, + const std::string& decode_method) { + ET_CHECK_OK_OR_RETURN_ERROR(load(prefill_method, decode_method)); + for (auto& buf : caches_) { + std::memset(buf.data, 0, buf.nbytes); + } + return Error::Ok; +} + +Result> HfStaticIOManager::prepare( + const TensorPtr& input, + const TensorPtr& start_pos, + const std::string& method) { + ET_CHECK_OK_OR_RETURN_ERROR(load_method(method)); + auto& io = methods_.at(method); + ET_CHECK_OR_RETURN_ERROR( + start_pos && input && start_pos->scalar_type() == ScalarType::Long && + start_pos->numel() == 1 && + (input->scalar_type() == ScalarType::Long || + input->scalar_type() == ScalarType::Int), + InvalidArgument, + "hf_static: expected int64 start_pos and int32/int64 tokens."); + const int64_t n = input->numel(); + const int64_t start = start_pos->const_data_ptr()[0]; + ET_CHECK_OR_RETURN_ERROR( + start >= 0, + InvalidArgument, + "hf_static: negative position %" PRId64 ".", + start); + ET_CHECK_OR_RETURN_ERROR( + n >= 1 && n <= io.query_len, + InvalidArgument, + "hf_static: %s takes at most %" PRId64 " tokens per step, got %" PRId64 + ".", + method.c_str(), + io.query_len, + n); + // The step writes all its rows, padding included. + ET_CHECK_OR_RETURN_ERROR( + start <= context_len_ - io.query_len, + OutOfResources, + "hf_static: %s at position %" PRId64 + " would write past the cache length %" PRId64 ".", + method.c_str(), + start, + context_len_); + // A short chunk is padded with its last token, at the following positions; + // the rows it writes past the real tokens are rewritten before they're read. + auto* ids = static_cast(io.step_inputs.at(io.input_ids).data); + auto* pos = static_cast(io.step_inputs.at(io.position_ids).data); + for (int64_t i = 0; i < io.query_len; ++i) { + const int64_t src = std::min(i, n - 1); + ids[i] = input->scalar_type() == ScalarType::Int + ? input->const_data_ptr()[src] + : input->const_data_ptr()[src]; + pos[i] = std::min(start + i, context_len_ - 1); + } + // Row i sees slots 0..start + i, as HF's generator builds the mask. A + // sliding window's (window-sized) mask also tells real rows from padding, + // which sees nothing: the program shifts the window by the real rows. + for (size_t m : io.masks) { + auto& buf = io.step_inputs.at(m); + const auto sizes = buf.tensor->sizes(); + const int64_t width = sizes[sizes.size() - 1]; + auto* mask = static_cast(buf.data); + for (int64_t i = 0; i < io.query_len; ++i) { + const bool padding = i >= n && width < context_len_; + const int64_t visible = padding ? 0 : std::min(width, start + i + 1); + std::memset(mask + i * width, 1, visible); + std::memset(mask + i * width + visible, 0, width - visible); + } + } + for (size_t c = 2; c < caches_.size(); c += 3) { + *static_cast(caches_[c].data) = start; + } + io.valid = n; + std::vector values; + values.reserve(io.num_inputs); + for (size_t i = 0; i < io.num_inputs; ++i) { + values.emplace_back( + i < io.first_cache ? *io.step_inputs.at(i).tensor + : *caches_[i - io.first_cache].tensor); + } + return values; +} + +Result> HfStaticIOManager::prepare_prefill( + const TensorPtr& input, + const TensorPtr& start_pos, + const std::string& prefill_method) { + return prepare(input, start_pos, prefill_method); +} + +Result> HfStaticIOManager::prepare_decode( + const TensorPtr& input, + const TensorPtr& start_pos, + const std::string& decode_method) { + return prepare(input, start_pos, decode_method); +} + +Error HfStaticIOManager::update_prefill( + const std::vector& model_outputs, + const std::string& prefill_method) { + (void)model_outputs; + (void)prefill_method; + return Error::Ok; +} + +Error HfStaticIOManager::update_decode( + const std::vector& model_outputs, + const std::string& decode_method) { + (void)model_outputs; + (void)decode_method; + return Error::Ok; +} + +Result HfStaticIOManager::logits( + const std::vector& model_outputs, + const std::string& method) { + // The program returns its mutated inputs before its own outputs. + const size_t index = static_cast(3 * num_layers_); + ET_CHECK_OR_RETURN_ERROR( + model_outputs.size() > index && model_outputs[index].isTensor(), + InvalidState, + "hf_static: expected the logits at output %zu of %zu.", + index, + model_outputs.size()); + const auto& t = model_outputs[index].toTensor(); + const auto& io = methods_.at(method); + const size_t vocab = logits_f_.size(); + // The row of the last real token (a prompt chunk may be padded). + const size_t row = t.dim() == 3 && t.size(1) > 1 ? io.valid - 1 : 0; + switch (t.scalar_type()) { + case ScalarType::Float: { + const float* src = t.const_data_ptr() + row * vocab; + std::copy(src, src + vocab, logits_f_.begin()); + break; + } + case ScalarType::Half: { + const auto* src = + t.const_data_ptr() + row * vocab; + for (size_t i = 0; i < vocab; ++i) { + logits_f_[i] = static_cast(src[i]); + } + break; + } + default: + ET_LOG( + Error, + "hf_static: unsupported logits dtype %d", + (int)t.scalar_type()); + return Error::InvalidArgument; + } + return *logits_f_tensor_; +} + +void register_hf_static_io_manager() { + static const bool registered = register_io_manager( + HfStaticIOManager::kContract, [](ET_MODULE_NAMESPACE::Module& module) { + return std::make_unique(module); + }); + (void)registered; +} + +} // namespace llm +} // namespace extension +} // namespace executorch diff --git a/extension/llm/runner/io_manager/hf_static_io_manager.h b/extension/llm/runner/io_manager/hf_static_io_manager.h new file mode 100644 index 00000000000..336f32d715c --- /dev/null +++ b/extension/llm/runner/io_manager/hf_static_io_manager.h @@ -0,0 +1,142 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include +#include + +#include + +namespace executorch { +namespace extension { +namespace llm { + +/** + * @brief Drives Hugging Face's static-cache ABI ("hf_static"). + * + * Each method is HF's decode graph at a static query length N. Inputs: + * input_ids [1, N], attention_mask(s) [1, 1, N, L] (bool), position_ids + * [1, N], then per layer keys and values [1, H, L, D] and cumulative_length. + * Outputs: the caches and counters updated in place, logits [1, N, V], and + * optionally the duplicate HF user cache state. + * A program may carry `prefill` (N tokens per step, for the prompt) and + * `decode` (one token), or a single `forward`. + * + * Version 1 supports float32 and float16 caches. See the versioned contract + * in docs/source/llm-hf-static-io-contract.md for metadata and padding rules. + * + * Sliding-window layers may hold window-sized caches (HF's): get_max_seq_len + * defines the context, and a window's mask also marks a prompt chunk's padding + * rows (all false), so the program shifts the window by the real rows only. + * + * This IOManager owns the caches and counters, one set shared by every method + * (from the I/O buffer allocator when there is one, so a delegate reads them in + * place), and builds each step's tokens, masks and positions as HF's generator + * does. The program writes each step's rows and counters into those buffers. + */ +class ET_EXPERIMENTAL HfStaticIOManager : public IOManager { + public: + static constexpr const char* kContract = "hf_static"; + static constexpr int64_t kContractVersion = 1; + + explicit HfStaticIOManager(ET_MODULE_NAMESPACE::Module& module); + ~HfStaticIOManager() override; + + ET_NODISCARD runtime::Error load( + const std::string& prefill_method, + const std::string& decode_method) override; + + ET_NODISCARD runtime::Error reset( + const std::string& prefill_method, + const std::string& decode_method) override; + + runtime::Result> prepare_prefill( + const TensorPtr& input, + const TensorPtr& start_pos, + const std::string& prefill_method) override; + + runtime::Result> prepare_decode( + const TensorPtr& input, + const TensorPtr& start_pos, + const std::string& decode_method) override; + + ET_NODISCARD runtime::Error update_prefill( + const std::vector& model_outputs, + const std::string& prefill_method) override; + + ET_NODISCARD runtime::Error update_decode( + const std::vector& model_outputs, + const std::string& decode_method) override; + + runtime::Result logits( + const std::vector& model_outputs, + const std::string& method) override; + + std::pair method_names( + const std::unordered_set& available) const override; + + int64_t max_tokens_per_step(const std::string& method) const override; + int64_t max_tokens_per_step(const std::string& method, int64_t start_pos) + const override; + + private: + struct Buffer { + void* data = nullptr; + size_t nbytes = 0; + TensorPtr tensor; + }; + // A method's per-step inputs; the caches are shared. + struct Method { + int64_t query_len = 1; + size_t num_inputs = 0; + // Input index -> buffer, for input_ids, the masks and position_ids. + std::map step_inputs; + size_t input_ids = 0; + size_t position_ids = 0; + std::vector masks; + size_t first_cache = 0; + // The logits, when the program leaves them to the runtime. + Buffer logits; + // The last step's real (unpadded) token count. + int64_t valid = 0; + }; + + runtime::Error load_metadata(); + runtime::Error load_method(const std::string& method); + runtime::Result> prepare( + const TensorPtr& input, + const TensorPtr& start_pos, + const std::string& method); + void allocate( + const std::vector>& buffers); + + ET_MODULE_NAMESPACE::Module& module_; + std::unique_ptr allocator_; + struct Block { + void* data; + bool from_allocator; + }; + std::vector blocks_; + std::map methods_; + // Per layer keys, values, cumulative length; shared by every method. + std::vector caches_; + bool metadata_loaded_ = false; + executorch::aten::ScalarType cache_dtype_ = + executorch::aten::ScalarType::Float; + int32_t num_layers_ = 0; + int64_t context_len_ = 0; + std::vector logits_f_; + TensorPtr logits_f_tensor_; +}; + +void register_hf_static_io_manager(); + +} // namespace llm +} // namespace extension +} // namespace executorch diff --git a/extension/llm/runner/io_manager/io_manager.h b/extension/llm/runner/io_manager/io_manager.h index fc9a8f0641b..44bf3cc3ab1 100644 --- a/extension/llm/runner/io_manager/io_manager.h +++ b/extension/llm/runner/io_manager/io_manager.h @@ -8,6 +8,13 @@ #pragma once +#include +#include +#include +#include +#include +#include + #include #include @@ -236,6 +243,74 @@ class ET_EXPERIMENTAL IOManager { return update_decode(model_outputs, "forward"); } + /** + * @brief The logits among a method's outputs, as the sampler reads them. + * + * The default takes a method that returns only its logits. An IOManager for + * a model that returns more (e.g. its updated cache) picks them out, and may + * hand back a converted copy (e.g. dequantized). + * + * @param model_outputs The outputs of the method just executed. + * @param method The method that produced them. + * @return The logits tensor. + */ + virtual runtime::Result logits( + const std::vector& model_outputs, + const std::string& method) { + (void)method; + if (model_outputs.size() != 1 || !model_outputs[0].isTensor()) { + ET_LOG( + Error, + "Expected the method to return only its logits tensor, got %zu outputs.", + model_outputs.size()); + return runtime::Error::InvalidState; + } + return model_outputs[0].toTensor(); + } + + /** + * @brief The most tokens one call of `method` takes, or -1 for no limit of + * the IOManager's own (e.g. a fixed-size prompt step). + */ + virtual int64_t max_tokens_per_step(const std::string& method) const { + (void)method; + return -1; + } + + virtual int64_t max_tokens_per_step( + const std::string& method, + int64_t start_pos) const { + (void)start_pos; + return max_tokens_per_step(method); + } + + /** + * @brief For an encoder-decoder contract: route `encoder_method`'s outputs + * into the decode step's fixed inputs. Call once after loading, then after + * each run of the encoder pass its outputs to `update_encoder`. + */ + ET_NODISCARD virtual runtime::Error load_encoder( + const std::string& encoder_method) { + (void)encoder_method; + return runtime::Error::NotSupported; + } + + ET_NODISCARD virtual runtime::Error update_encoder( + const std::vector& encoder_outputs) { + (void)encoder_outputs; + return runtime::Error::NotSupported; + } + + /** + * @brief The prefill and decode method names this IOManager drives, given + * the methods the program has. Empty strings leave the runner's defaults. + */ + virtual std::pair method_names( + const std::unordered_set& available) const { + (void)available; + return {"", ""}; + } + private: /** * @brief Reference to the Module used for method metadata and execution. @@ -243,6 +318,87 @@ class ET_EXPERIMENTAL IOManager { ET_MODULE_NAMESPACE::Module& module_; }; +using IOManagerFactory = + std::function(ET_MODULE_NAMESPACE::Module&)>; + +/** + * @brief IOManagers by the I/O contract they speak. + * + * A program names its contract in the `get_io_contract` constant method; the + * runner builds the IOManager registered for it, or the default IOManager when + * the program names none. + */ +inline std::unordered_map& +io_manager_registry() { + static std::unordered_map registry; + return registry; +} + +inline bool register_io_manager( + const std::string& contract, + IOManagerFactory factory) { + return io_manager_registry().emplace(contract, std::move(factory)).second; +} + +inline runtime::Result> make_io_manager( + ET_MODULE_NAMESPACE::Module& module, + const std::string& contract) { + if (!contract.empty()) { + auto it = io_manager_registry().find(contract); + if (it != io_manager_registry().end()) { + auto manager = it->second(module); + if (!manager) { + return runtime::Error::InvalidState; + } + return manager; + } + ET_LOG( + Error, + "No IOManager registered for I/O contract '%s'.", + contract.c_str()); + return runtime::Error::NotSupported; + } + return std::make_unique(module); +} + +/// Registers the IOManagers this library implements. Idempotent. +void register_builtin_io_managers(); + +/** + * @brief Memory for the I/O buffers an IOManager owns. + * + * A backend whose device reads graph I/O from memory it can map (for example + * QNN's rpcmem) registers an allocator, so the device uses these buffers in + * place instead of copying them on every call. Without one, buffers come from + * the heap. + */ +class ET_EXPERIMENTAL IOBufferAllocator { + public: + virtual ~IOBufferAllocator() = default; + /// A block of at least `nbytes`, or nullptr to fall back to the heap. + virtual void* allocate(size_t nbytes, size_t alignment) = 0; + virtual void deallocate(void* block) = 0; + /// Declares a tensor whose data starts at `tensor_data` inside `block`. + virtual void bind(void* block, void* tensor_data) { + (void)block; + (void)tensor_data; + } +}; + +using IOBufferAllocatorFactory = + std::function()>; + +inline IOBufferAllocatorFactory& io_buffer_allocator_factory() { + static IOBufferAllocatorFactory factory; + return factory; +} + +/// Sets the allocator IOManagers use for their I/O buffers. The last call wins. +inline bool set_io_buffer_allocator(IOBufferAllocatorFactory factory) { + io_buffer_allocator_factory() = std::move(factory); + return true; +} + } // namespace llm } // namespace extension } // namespace executorch diff --git a/extension/llm/runner/io_manager/kv_cache_io_manager.cpp b/extension/llm/runner/io_manager/kv_cache_io_manager.cpp new file mode 100644 index 00000000000..5ded44f79ad --- /dev/null +++ b/extension/llm/runner/io_manager/kv_cache_io_manager.cpp @@ -0,0 +1,555 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include + +#include +#include +#include +#include + +#include + +namespace executorch { +namespace extension { +namespace llm { + +using ::executorch::aten::ScalarType; +using ::executorch::runtime::Error; +using ::executorch::runtime::EValue; +using ::executorch::runtime::Result; + +namespace { + +constexpr const char* kPromptMethod = "prefill_forward"; +constexpr const char* kStepMethod = "kv_forward"; +constexpr size_t kBufferAlignment = 64; + +size_t align_up(size_t n) { + return (n + kBufferAlignment - 1) / kBufferAlignment * kBufferAlignment; +} + +size_t element_size(ScalarType dtype) { + return ::executorch::runtime::elementSize(dtype); +} + +template +void fill(std::byte* base, size_t offset, size_t count, T value) { + std::fill_n(reinterpret_cast(base) + offset, count, value); +} + +void fill_value( + ScalarType dtype, + std::byte* base, + size_t offset, + size_t count, + bool visible) { + switch (dtype) { + case ScalarType::Float: + fill(base, offset, count, visible ? 0.0f : -65535.0f); + break; + case ScalarType::UInt16: + fill(base, offset, count, visible ? 65535u : 0u); + break; + case ScalarType::Byte: + fill(base, offset, count, visible ? 255u : 0u); + break; + default: + ET_CHECK_MSG(false, "Unsupported attention mask dtype %d", (int)dtype); + } +} + +void write_index(ScalarType dtype, std::byte* base, size_t i, int64_t value) { + if (dtype == ScalarType::Int) { + reinterpret_cast(base)[i] = static_cast(value); + } else { + reinterpret_cast(base)[i] = value; + } +} + +} // namespace + +KvCacheIOManager::KvCacheIOManager(ET_MODULE_NAMESPACE::Module& module) + : IOManager(module), module_(module) { + if (auto& factory = io_buffer_allocator_factory()) { + allocator_ = factory(); + } +} + +KvCacheIOManager::~KvCacheIOManager() { + methods_.clear(); + past_k_.clear(); + past_v_.clear(); + for (const auto& block : blocks_) { + if (block.from_allocator) { + allocator_->deallocate(block.data); + } else { + ::operator delete(block.data, std::align_val_t(kBufferAlignment)); + } + } +} + +void KvCacheIOManager::allocate( + const std::vector>& + buffers) { + size_t total = 0; + for (const auto& [buf, info] : buffers) { + total = align_up(total) + info.nbytes(); + } + void* block = + allocator_ ? allocator_->allocate(total, kBufferAlignment) : nullptr; + const bool from_allocator = block != nullptr; + if (!from_allocator) { + block = ::operator new(total, std::align_val_t(kBufferAlignment)); + } + blocks_.push_back({block, from_allocator}); + std::memset(block, 0, total); + size_t offset = 0; + for (const auto& [buf, info] : buffers) { + offset = align_up(offset); + buf->data = static_cast(block) + offset; + buf->nbytes = info.nbytes(); + std::vector sizes( + info.sizes().begin(), info.sizes().end()); + buf->tensor = from_blob(buf->data, sizes, info.scalar_type()); + if (from_allocator) { + allocator_->bind(block, buf->data); + } + offset += info.nbytes(); + } +} + +std::pair KvCacheIOManager::method_names( + const std::unordered_set& available) const { + if (!available.count(kStepMethod)) { + return {"", ""}; + } + // A separate prompt graph takes the prompt in fixed-size chunks; without + // one, the decode step serves the prompt too. + return { + available.count(kPromptMethod) ? kPromptMethod : kStepMethod, + kStepMethod}; +} + +int64_t KvCacheIOManager::max_tokens_per_step(const std::string& method) const { + auto it = methods_.find(method); + return it == methods_.end() ? -1 : it->second.ar_len; +} + +void KvCacheIOManager::rearrange(int32_t past_len) { + // Move the filled positions into the layout of `past_len`: K rows and V + // heads are strided by the past length, so each moves separately. + if (active_past_len_ >= 0 && active_past_len_ != past_len && filled_ > 0) { + const int32_t src_len = active_past_len_; + for (int32_t l = 0; l < num_layers_; ++l) { + auto& ks = past_k_[src_len][l]; + auto& kd = past_k_[past_len][l]; + const size_t es = element_size(ks.tensor->scalar_type()); + const size_t rows = ks.tensor->numel() / src_len; + for (size_t r = 0; r < rows; ++r) { + std::memcpy( + kd.data + r * past_len * es, + ks.data + r * src_len * es, + filled_ * es); + } + auto& vs = past_v_[src_len][l]; + auto& vd = past_v_[past_len][l]; + const size_t heads = vs.tensor->size(1); + const size_t dim = vs.tensor->size(3); + for (size_t h = 0; h < heads; ++h) { + std::memcpy( + vd.data + h * past_len * dim * es, + vs.data + h * src_len * dim * es, + filled_ * dim * es); + } + } + } + active_past_len_ = past_len; +} + +Error KvCacheIOManager::load( + const std::string& prefill_method, + const std::string& decode_method) { + // Quantized logits are dequantized with the scale the program carries, if + // any; without it they keep their order, which is all greedy decoding needs. + auto names = module_.method_names(); + if (names.ok() && names->count("get_logits_scale") && + names->count("get_logits_zero_point")) { + logits_scale_ = module_.get("get_logits_scale")->toDouble(); + logits_zero_point_ = module_.get("get_logits_zero_point")->toInt(); + } + if (names.ok() && names->count("get_n_cross_layers")) { + num_cross_layers_ = + static_cast(module_.get("get_n_cross_layers")->toInt()); + } + // The runner may also call load() with its default method names; only the + // program's own methods are this contract's. + for (const auto& method : {prefill_method, decode_method}) { + if (names.ok() && names->count(method)) { + ET_CHECK_OK_OR_RETURN_ERROR(load_method(method)); + } + } + return Error::Ok; +} + +Error KvCacheIOManager::load_method(const std::string& method) { + if (methods_.count(method)) { + return Error::Ok; + } + ET_CHECK_OK_OR_RETURN_ERROR(module_.load_method(method)); + auto meta_result = module_.method_meta(method); + ET_CHECK_OK_OR_RETURN_ERROR(meta_result.error()); + auto meta = std::move(*meta_result); + const size_t num_inputs = meta.num_inputs(); + const size_t cross_inputs = 2 * static_cast(num_cross_layers_); + ET_CHECK_OR_RETURN_ERROR( + num_inputs >= 5 + cross_inputs && + (num_inputs - 3 - cross_inputs) % 2 == 0, + InvalidArgument, + "%s has %zu inputs; a kvcache_io method takes tokens, mask, positions, a K and V cache per layer and %zu cross-attention inputs.", + method.c_str(), + num_inputs, + cross_inputs); + const int32_t layers = + static_cast((num_inputs - 3 - cross_inputs) / 2); + ET_CHECK_OR_RETURN_ERROR( + num_layers_ == 0 || num_layers_ == layers, + InvalidArgument, + "%s has %d layers, another method %d.", + method.c_str(), + layers, + num_layers_); + num_layers_ = layers; + + MethodIO io; + std::vector> buffers; + auto make = [&buffers]( + Buffer& buf, const ::executorch::runtime::TensorInfo& info) { + buffers.emplace_back(&buf, info); + }; + auto tokens = *meta.input_tensor_meta(0); + auto mask = *meta.input_tensor_meta(1); + io.ar_len = tokens.sizes()[tokens.sizes().size() - 1]; + io.context_len = mask.sizes()[mask.sizes().size() - 1]; + io.past_len = io.context_len - io.ar_len; + io.mask_dtype = mask.scalar_type(); + make(io.tokens, tokens); + make(io.mask, mask); + make(io.positions, *meta.input_tensor_meta(2)); + + auto& past_k = past_k_[io.past_len]; + auto& past_v = past_v_[io.past_len]; + const bool fresh = past_k.empty(); + if (fresh) { + past_k.resize(layers); + past_v.resize(layers); + } + for (int32_t l = 0; l < layers; ++l) { + if (fresh) { + make(past_k[l], *meta.input_tensor_meta(3 + l)); + make(past_v[l], *meta.input_tensor_meta(3 + layers + l)); + } + io.past_k.push_back(&past_k[l]); + io.past_v.push_back(&past_v[l]); + } + + make(io.logits, *meta.output_tensor_meta(0)); + io.new_k.resize(layers); + io.new_v.resize(layers); + for (int32_t l = 0; l < layers; ++l) { + make(io.new_k[l], *meta.output_tensor_meta(1 + l)); + make(io.new_v[l], *meta.output_tensor_meta(1 + layers + l)); + } + if (cross_.empty() && num_cross_layers_ > 0) { + cross_.resize(cross_inputs); + for (size_t i = 0; i < cross_inputs; ++i) { + make(cross_[i], *meta.input_tensor_meta(3 + 2 * layers + i)); + } + } + allocate(buffers); + // Outputs left out of memory planning are written into these buffers; + // planned ones (e.g. a CPU program) come back in the program's own memory. + auto bind = [&](const Buffer& buf, size_t index) -> Error { + if (meta.output_tensor_meta(index)->is_memory_planned()) { + return Error::Ok; + } + return module_.set_output(method, EValue(*buf.tensor), index); + }; + ET_CHECK_OK_OR_RETURN_ERROR(bind(io.logits, 0)); + for (int32_t l = 0; l < layers; ++l) { + ET_CHECK_OK_OR_RETURN_ERROR(bind(io.new_k[l], 1 + l)); + ET_CHECK_OK_OR_RETURN_ERROR(bind(io.new_v[l], 1 + layers + l)); + } + const auto vocab = meta.output_tensor_meta(0)->sizes(); + logits_f_.assign(vocab[vocab.size() - 1], 0.0f); + logits_f_tensor_ = from_blob( + logits_f_.data(), + {1, 1, static_cast(logits_f_.size())}, + ScalarType::Float); + methods_.emplace(method, std::move(io)); + return Error::Ok; +} + +Error KvCacheIOManager::reset( + const std::string& prefill_method, + const std::string& decode_method) { + (void)prefill_method; + (void)decode_method; + for (auto* caches : {&past_k_, &past_v_}) { + for (auto& [len, layers] : *caches) { + for (auto& buf : layers) { + std::memset(buf.data, 0, buf.nbytes); + } + } + } + filled_ = 0; + active_past_len_ = -1; + return Error::Ok; +} + +void KvCacheIOManager::fill_mask( + MethodIO& io, + int64_t n_past, + int32_t n_valid) { + std::byte* base = io.mask.data; + const size_t row = io.context_len; + fill_value(io.mask_dtype, base, 0, row * io.ar_len, /*visible=*/false); + for (int32_t r = 0; r < io.ar_len; ++r) { + // Every token sees the filled past and the step's tokens up to itself. + fill_value(io.mask_dtype, base, r * row, n_past, true); + const int32_t self = std::min(r, std::max(n_valid - 1, 0)); + fill_value(io.mask_dtype, base, r * row + io.past_len, self + 1, true); + } +} + +Result> KvCacheIOManager::prepare( + const TensorPtr& input, + const TensorPtr& start_pos, + const std::string& method) { + ET_CHECK_OK_OR_RETURN_ERROR(load_method(method)); + auto& io = methods_.at(method); + const int32_t n = static_cast(input->numel()); + const int64_t start = start_pos->const_data_ptr()[0]; + ET_CHECK_OR_RETURN_ERROR( + n >= 1 && n <= io.ar_len, + InvalidArgument, + "%s takes at most %d tokens per step, got %d.", + method.c_str(), + io.ar_len, + n); + ET_CHECK_OR_RETURN_ERROR( + start + n <= io.past_len + (io.ar_len == 1 ? 1 : 0), + OutOfResources, + "Context full: position %" PRId64 " with %d slots of past.", + start, + io.past_len); + const auto token_dtype = io.tokens.tensor->scalar_type(); + const auto pos_dtype = io.positions.tensor->scalar_type(); + for (int32_t i = 0; i < io.ar_len; ++i) { + const int32_t src = std::min(i, n - 1); + const int64_t token = input->scalar_type() == ScalarType::Int + ? input->const_data_ptr()[src] + : input->const_data_ptr()[src]; + write_index(token_dtype, io.tokens.data, i, token); + write_index( + pos_dtype, + io.positions.data, + i, + std::min(start + i, io.context_len - 1)); + } + rearrange(io.past_len); + fill_mask(io, start, n); + io.valid = n; + io.start = start; + + std::vector inputs; + inputs.reserve(3 + 2 * num_layers_ + cross_.size()); + inputs.emplace_back(*io.tokens.tensor); + inputs.emplace_back(*io.mask.tensor); + inputs.emplace_back(*io.positions.tensor); + for (auto* buf : io.past_k) { + inputs.emplace_back(*buf->tensor); + } + for (auto* buf : io.past_v) { + inputs.emplace_back(*buf->tensor); + } + for (auto& buf : cross_) { + inputs.emplace_back(*buf.tensor); + } + return inputs; +} + +Error KvCacheIOManager::update( + const std::vector& model_outputs, + const std::string& method) { + auto& io = methods_.at(method); + ET_CHECK_OR_RETURN_ERROR( + model_outputs.size() == static_cast(1 + 2 * num_layers_), + InvalidState, + "%s returned %zu outputs, expected %d.", + method.c_str(), + model_outputs.size(), + 1 + 2 * num_layers_); + for (int32_t l = 0; l < num_layers_; ++l) { + // K is [1, H, D, ar] -> past [1, H, D, P]: one strided write per (head, + // dim) row. + const auto& k_new = model_outputs[1 + l].toTensor(); + auto& k_past = *io.past_k[l]; + const size_t es = element_size(k_new.scalar_type()); + const size_t rows = k_new.numel() / io.ar_len; + const std::byte* src = k_new.const_data_ptr(); + std::byte* dst = k_past.data; + for (size_t r = 0; r < rows; ++r) { + std::memcpy( + dst + (r * io.past_len + io.start) * es, + src + r * io.ar_len * es, + io.valid * es); + } + // V is [1, H, ar, D] -> past [1, H, P, D]: one contiguous write per head. + const auto& v_new = model_outputs[1 + num_layers_ + l].toTensor(); + auto& v_past = *io.past_v[l]; + const size_t heads = v_new.size(1); + const size_t dim = v_new.size(3); + const std::byte* vsrc = v_new.const_data_ptr(); + std::byte* vdst = v_past.data; + for (size_t h = 0; h < heads; ++h) { + std::memcpy( + vdst + (h * io.past_len + io.start) * dim * es, + vsrc + h * io.ar_len * dim * es, + io.valid * dim * es); + } + } + filled_ = io.start + io.valid; + return Error::Ok; +} + +Error KvCacheIOManager::load_encoder(const std::string& encoder_method) { + ET_CHECK_OR_RETURN_ERROR( + !cross_.empty(), + InvalidState, + "Load the decode method of a program with cross-attention inputs first."); + ET_CHECK_OK_OR_RETURN_ERROR(module_.load_method(encoder_method)); + auto meta = module_.method_meta(encoder_method); + ET_CHECK_OK_OR_RETURN_ERROR(meta.error()); + ET_CHECK_OR_RETURN_ERROR( + meta->num_outputs() == cross_.size(), + InvalidArgument, + "%s returns %zu outputs; the decode step takes %zu cross-attention inputs.", + encoder_method.c_str(), + meta->num_outputs(), + cross_.size()); + // The encoder writes straight into the decode step's inputs where it can. + for (size_t i = 0; i < cross_.size(); ++i) { + if (!meta->output_tensor_meta(i)->is_memory_planned()) { + ET_CHECK_OK_OR_RETURN_ERROR( + module_.set_output(encoder_method, EValue(*cross_[i].tensor), i)); + } + } + return Error::Ok; +} + +Error KvCacheIOManager::update_encoder( + const std::vector& encoder_outputs) { + ET_CHECK_OR_RETURN_ERROR( + encoder_outputs.size() == cross_.size(), + InvalidArgument, + "Got %zu encoder outputs for %zu cross-attention inputs.", + encoder_outputs.size(), + cross_.size()); + for (size_t i = 0; i < cross_.size(); ++i) { + const auto& out = encoder_outputs[i].toTensor(); + if (out.const_data_ptr() != cross_[i].data) { + ET_CHECK_OR_RETURN_ERROR( + out.nbytes() == cross_[i].nbytes, + InvalidArgument, + "Encoder output %zu has %zu bytes, the input %zu.", + i, + out.nbytes(), + cross_[i].nbytes); + std::memcpy(cross_[i].data, out.const_data_ptr(), out.nbytes()); + } + } + return Error::Ok; +} + +Result> KvCacheIOManager::prepare_prefill( + const TensorPtr& input, + const TensorPtr& start_pos, + const std::string& prefill_method) { + return prepare(input, start_pos, prefill_method); +} + +Result> KvCacheIOManager::prepare_decode( + const TensorPtr& input, + const TensorPtr& start_pos, + const std::string& decode_method) { + return prepare(input, start_pos, decode_method); +} + +Error KvCacheIOManager::update_prefill( + const std::vector& model_outputs, + const std::string& prefill_method) { + return update(model_outputs, prefill_method); +} + +Error KvCacheIOManager::update_decode( + const std::vector& model_outputs, + const std::string& decode_method) { + return update(model_outputs, decode_method); +} + +Result KvCacheIOManager::logits( + const std::vector& model_outputs, + const std::string& method) { + auto& io = methods_.at(method); + const auto& t = model_outputs[0].toTensor(); + const size_t vocab = logits_f_.size(); + // The row of the last real token (a prompt chunk may be padded). + const size_t row = t.size(1) > 1 ? static_cast(io.valid - 1) : 0; + switch (t.scalar_type()) { + case ScalarType::Float: { + const float* src = t.const_data_ptr() + row * vocab; + std::copy(src, src + vocab, logits_f_.begin()); + break; + } + case ScalarType::UInt16: { + const uint16_t* src = t.const_data_ptr() + row * vocab; + for (size_t i = 0; i < vocab; ++i) { + logits_f_[i] = static_cast( + (static_cast(src[i]) - logits_zero_point_) * + logits_scale_); + } + break; + } + case ScalarType::Byte: { + const uint8_t* src = t.const_data_ptr() + row * vocab; + for (size_t i = 0; i < vocab; ++i) { + logits_f_[i] = static_cast( + (static_cast(src[i]) - logits_zero_point_) * + logits_scale_); + } + break; + } + default: + ET_LOG(Error, "Unsupported logits dtype %d", (int)t.scalar_type()); + return Error::InvalidArgument; + } + return *logits_f_tensor_; +} + +void register_kv_cache_io_manager() { + static const bool registered = register_io_manager( + KvCacheIOManager::kContract, [](ET_MODULE_NAMESPACE::Module& module) { + return std::make_unique(module); + }); + (void)registered; +} + +} // namespace llm +} // namespace extension +} // namespace executorch diff --git a/extension/llm/runner/io_manager/kv_cache_io_manager.h b/extension/llm/runner/io_manager/kv_cache_io_manager.h new file mode 100644 index 00000000000..2efd7e23330 --- /dev/null +++ b/extension/llm/runner/io_manager/kv_cache_io_manager.h @@ -0,0 +1,154 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#pragma once + +#include +#include +#include +#include +#include + +#include + +namespace executorch { +namespace extension { +namespace llm { + +/** + * @brief Drives a decoder step that takes its KV cache as I/O ("kvcache_io"). + * + * The program's method is + * + * (input_tokens, atten_mask, pos_ids, past_k_0..N-1, past_v_0..N-1) + * -> (logits, new_k_0..N-1, new_v_0..N-1) + * + * with `past_k_i` stored transposed `[1, H, D, P]`, `past_v_i` `[1, H, P, D]`, + * (an encoder-decoder's step also takes `cross_k_i` / `cross_v_i`, fixed per + * input and written by its encoder method, after the past caches), + * only the new tokens' K/V slice returned, and an additive (or, quantized, + * 0/max) attention mask over the `P` past slots followed by the step's own + * slots. This IOManager owns the cache and mask buffers: before each step it + * opens the mask over the filled past slots; after it, it writes the new slice + * into the past buffers at the current position. The runner above only sees + * tokens in and logits out. + */ +class ET_EXPERIMENTAL KvCacheIOManager : public IOManager { + public: + static constexpr const char* kContract = "kvcache_io"; + + explicit KvCacheIOManager(ET_MODULE_NAMESPACE::Module& module); + ~KvCacheIOManager() override; + + ET_NODISCARD runtime::Error load( + const std::string& prefill_method, + const std::string& decode_method) override; + + ET_NODISCARD runtime::Error reset( + const std::string& prefill_method, + const std::string& decode_method) override; + + runtime::Result> prepare_prefill( + const TensorPtr& input, + const TensorPtr& start_pos, + const std::string& prefill_method) override; + + runtime::Result> prepare_decode( + const TensorPtr& input, + const TensorPtr& start_pos, + const std::string& decode_method) override; + + ET_NODISCARD runtime::Error update_prefill( + const std::vector& model_outputs, + const std::string& prefill_method) override; + + ET_NODISCARD runtime::Error update_decode( + const std::vector& model_outputs, + const std::string& decode_method) override; + + runtime::Result logits( + const std::vector& model_outputs, + const std::string& method) override; + + std::pair method_names( + const std::unordered_set& available) const override; + + int64_t max_tokens_per_step(const std::string& method) const override; + + ET_NODISCARD runtime::Error load_encoder( + const std::string& encoder_method) override; + + ET_NODISCARD runtime::Error update_encoder( + const std::vector& encoder_outputs) override; + + private: + struct Buffer { + std::byte* data = nullptr; + size_t nbytes = 0; + TensorPtr tensor; + }; + struct Block { + void* data; + bool from_allocator; + }; + + // One method's I/O. Cache buffers are shared by every method with the same + // past length. + struct MethodIO { + int32_t ar_len = 1; + int32_t past_len = 0; + int32_t context_len = 0; + Buffer tokens, mask, positions, logits; + std::vector new_k, new_v; + std::vector past_k, past_v; + executorch::aten::ScalarType mask_dtype; + int32_t valid = 0; // tokens of the last step that were real (not padding) + int64_t start = 0; + }; + + runtime::Error load_method(const std::string& method); + runtime::Result> prepare( + const TensorPtr& input, + const TensorPtr& start_pos, + const std::string& method); + runtime::Error update( + const std::vector& model_outputs, + const std::string& method); + // Places each buffer in one block from the I/O buffer allocator (the heap if + // there is none or it declines). + void allocate( + const std::vector>& + buffers); + void fill_mask(MethodIO& io, int64_t n_past, int32_t n_valid); + void rearrange(int32_t past_len); + + ET_MODULE_NAMESPACE::Module& module_; + std::unique_ptr allocator_; + std::vector blocks_; + std::unordered_map methods_; + // Past caches by past length, then layer. + std::unordered_map> past_k_, past_v_; + int32_t num_layers_ = 0; + int32_t num_cross_layers_ = 0; + // Cross-attention keys then values, shared by every method. + std::vector cross_; + // Positions filled in the cache, and the past length of the layout that + // holds them (methods with different step sizes have different layouts). + int64_t filled_ = 0; + int32_t active_past_len_ = -1; + std::vector logits_f_; + TensorPtr logits_f_tensor_; + double logits_scale_ = 1.0; + int64_t logits_zero_point_ = 0; +}; + +void register_kv_cache_io_manager(); + +} // namespace llm +} // namespace extension +} // namespace executorch diff --git a/extension/llm/runner/io_manager/targets.bzl b/extension/llm/runner/io_manager/targets.bzl index fe7fe9d56ae..081c8b15afb 100644 --- a/extension/llm/runner/io_manager/targets.bzl +++ b/extension/llm/runner/io_manager/targets.bzl @@ -17,3 +17,29 @@ def define_common_targets(): ], visibility = ["PUBLIC"], ) + + runtime.cxx_library( + name = "hf_static_io_manager" + aten_suffix, + exported_headers = ["hf_static_io_manager.h"], + srcs = ["hf_static_io_manager.cpp"], + exported_deps = [ + ":io_manager" + aten_suffix, + "//executorch/extension/llm/runner:constants", + ], + visibility = ["PUBLIC"], + ) + + # IOManager for the "kvcache_io" contract: KV cache as method I/O. + runtime.cxx_library( + name = "kv_cache_io_manager" + aten_suffix, + exported_headers = [ + "kv_cache_io_manager.h", + ], + srcs = [ + "kv_cache_io_manager.cpp", + ], + exported_deps = [ + ":io_manager" + aten_suffix, + ], + visibility = ["PUBLIC"], + ) diff --git a/extension/llm/runner/llm_runner_helper.cpp b/extension/llm/runner/llm_runner_helper.cpp index 105b3f0c9f4..bad6efb4c1f 100644 --- a/extension/llm/runner/llm_runner_helper.cpp +++ b/extension/llm/runner/llm_runner_helper.cpp @@ -10,6 +10,9 @@ #include #include + +#include +#include #include #include #include @@ -26,11 +29,18 @@ #include #include +#include + namespace executorch::extension::llm { using ::executorch::extension::Module; using ::executorch::runtime::Error; +void register_builtin_io_managers() { + register_kv_cache_io_manager(); + register_hf_static_io_manager(); +} + std::unique_ptr load_tokenizer( const std::string& tokenizer_path, std::unique_ptr> special_tokens, @@ -126,6 +136,13 @@ get_llm_metadata(tokenizers::Tokenizer* tokenizer, Module* module) { if (!get_result.ok()) { return get_result.error(); } + if (!get_result->isInt() && !get_result->isBool()) { + ET_LOG( + Error, + "Metadata method %s must return an integer or boolean.", + method_name.c_str()); + return Error::InvalidArgument; + } value = get_result->toScalar().to(); } else { ET_LOG( @@ -148,6 +165,15 @@ get_llm_metadata(tokenizers::Tokenizer* tokenizer, Module* module) { metadata[llm::kMaxContextLen]); } + for (const auto* name : {kMaxSeqLen, kMaxContextLen}) { + const auto value = metadata.at(name); + ET_CHECK_OR_RETURN_ERROR( + value > 0 && value <= std::numeric_limits::max(), + InvalidArgument, + "Metadata %s must be in [1, INT32_MAX].", + name); + } + // Set tokenizer-related metadata metadata[llm::kBosId] = tokenizer->bos_tok(); metadata[llm::kVocabSize] = tokenizer->vocab_size(); @@ -266,8 +292,32 @@ std::unique_ptr create_text_llm_runner( auto eos_ids = std::make_unique>( llm::get_eos_ids(tokenizer.get(), module.get())); - // Create IOManager - std::unique_ptr io_manager = std::make_unique(*module); + // Create the IOManager for the I/O contract the program declares, if any. + register_builtin_io_managers(); + std::string io_contract; + if (auto names = module->method_names(); + names.ok() && names->count(kIoContract)) { + auto contract = module->get(kIoContract); + if (!contract.ok() || !contract->isString()) { + ET_LOG(Error, "Failed to read the program's I/O contract label."); + return nullptr; + } + io_contract = std::string(contract->toString()); + if (io_contract.empty()) { + ET_LOG(Error, "The program's I/O contract label is empty."); + return nullptr; + } + ET_LOG(Info, "I/O contract: %s", io_contract.c_str()); + } + auto io_manager_result = make_io_manager(*module, io_contract); + if (!io_manager_result.ok()) { + ET_LOG( + Error, + "Failed to create the IOManager: 0x%x.", + static_cast(io_manager_result.error())); + return nullptr; + } + auto io_manager = std::move(*io_manager_result); // Read vocab_size for Sampler int32_t vocab_size = static_cast(metadata.at(kVocabSize)); @@ -289,6 +339,17 @@ std::unique_ptr create_text_llm_runner( method_names.count(kPrefillMethod) > 0) { prefill_method_name = kPrefillMethod; } + // An IOManager may drive methods of its own contract's names. + auto [io_prefill, io_decode] = io_manager->method_names(method_names); + if (!io_prefill.empty()) { + prefill_method_name = io_prefill; + } + if (io_manager->load( + prefill_method_name, + io_decode.empty() ? prefill_method_name : io_decode) != Error::Ok) { + ET_LOG(Error, "Failed to load the IOManager"); + return nullptr; + } // Create text_decoder_runner ET_LOG(Info, "Using method: %s", prefill_method_name.c_str()); @@ -302,7 +363,15 @@ std::unique_ptr create_text_llm_runner( // The decode stage gets its own decoder runner when the PTE exports a // separate decode method; otherwise both stages share text_decoder_runner. std::unique_ptr decode_text_decoder_runner; - if (method_names.count(kDecodeMethod) > 0 && + if (!io_decode.empty() && io_decode != prefill_method_name) { + decode_text_decoder_runner = std::make_unique( + module.get(), + io_manager.get(), + io_decode, + std::make_unique(vocab_size, init_temp), + stats.get()); + } else if ( + io_decode.empty() && method_names.count(kDecodeMethod) > 0 && prefill_method_name != kDecodeMethod) { ET_LOG(Info, "Using decode method: %s", kDecodeMethod); decode_text_decoder_runner = std::make_unique( @@ -316,12 +385,17 @@ std::unique_ptr create_text_llm_runner( ? decode_text_decoder_runner.get() : text_decoder_runner.get(); - // Create text_prefiller + // Create text_prefiller. A prompt method with a fixed step size takes the + // prompt in chunks of that size. + const int64_t step_tokens = + io_manager->max_tokens_per_step(prefill_method_name); auto text_prefiller = std::make_unique( text_decoder_runner.get(), metadata.at(kUseKVCache), - metadata.at(kEnableDynamicShape), - metadata.at(kMaxSeqLen)); + step_tokens > 1 ? true + : static_cast(metadata.at(kEnableDynamicShape)), + step_tokens > 1 ? step_tokens : metadata.at(kMaxSeqLen), + step_tokens > 1 ? decode_runner : nullptr); // Create text_token_generator with stats auto text_token_generator = std::make_unique( diff --git a/extension/llm/runner/targets.bzl b/extension/llm/runner/targets.bzl index 2cecf9975f2..ca032ad44a0 100644 --- a/extension/llm/runner/targets.bzl +++ b/extension/llm/runner/targets.bzl @@ -161,6 +161,8 @@ def define_common_targets(): ":text_prefiller" + aten_suffix, ":text_token_generator" + aten_suffix, "//executorch/extension/llm/runner/io_manager:io_manager" + aten_suffix, + "//executorch/extension/llm/runner/io_manager:hf_static_io_manager" + aten_suffix, + "//executorch/extension/llm/runner/io_manager:kv_cache_io_manager" + aten_suffix, "//executorch/extension/memory_allocator:cpu_caching_allocator", "//pytorch/tokenizers:hf_tokenizer", "//pytorch/tokenizers:llama2c_tokenizer", diff --git a/extension/llm/runner/test/BUCK b/extension/llm/runner/test/BUCK index 9ed85acadb7..cf0f2ea0136 100644 --- a/extension/llm/runner/test/BUCK +++ b/extension/llm/runner/test/BUCK @@ -45,3 +45,50 @@ fbcode_target(_kind = runtime.cxx_test, "NO_KVCACHE": "$(location fbcode//executorch/test/models:exported_programs[ModuleNoKVCache.pte])", } ) + +fbcode_target(_kind = runtime.python_library, + name = "export_hf_static_program_lib", + srcs = ["export_hf_static_program.py"], + base_module = "executorch.extension.llm.runner.test", + deps = [ + "//caffe2:torch", + "//executorch/exir:lib", + ], +) + +fbcode_target(_kind = runtime.python_binary, + name = "export_hf_static_program", + main_module = "executorch.extension.llm.runner.test.export_hf_static_program", + deps = [":export_hf_static_program_lib"], +) + +# A tiny program with Hugging Face's static-cache ABI. +fbcode_target(_kind = runtime.genrule, + name = "hf_static_program", + cmd = "$(exe :export_hf_static_program) --outdir $OUT", + outs = { + "HfStaticTiny.pte": ["HfStaticTiny.pte"], + "HfStaticTinyHalf.pte": ["HfStaticTinyHalf.pte"], + "HfStaticTinyWithState.pte": ["HfStaticTinyWithState.pte"], + }, + default_outs = ["."], +) + +fbcode_target(_kind = runtime.cxx_test, + name = "test_hf_static_io_manager", + srcs = ["test_hf_static_io_manager.cpp"], + deps = [ + "//executorch/extension/llm/runner:runner_lib", + "//executorch/extension/llm/runner/io_manager:hf_static_io_manager", + "//executorch/extension/llm/runner/io_manager:io_manager", + "//executorch/extension/module:module", + "//executorch/extension/tensor:tensor", + "//executorch/kernels/portable:generated_lib", + ], + env = { + "ET_HF_STATIC_FIXTURES_DIR": "$(location :hf_static_program)", + "ET_HF_STATIC_TINY_PATH": "$(location :hf_static_program[HfStaticTiny.pte])", + "ET_HF_STATIC_TINY_HALF_PATH": "$(location :hf_static_program[HfStaticTinyHalf.pte])", + "ET_HF_STATIC_TINY_STATE_PATH": "$(location :hf_static_program[HfStaticTinyWithState.pte])", + }, +) diff --git a/extension/llm/runner/test/CMakeLists.txt b/extension/llm/runner/test/CMakeLists.txt index cd380429601..f896abf126e 100644 --- a/extension/llm/runner/test/CMakeLists.txt +++ b/extension/llm/runner/test/CMakeLists.txt @@ -19,6 +19,8 @@ include(${EXECUTORCH_ROOT}/tools/cmake/Test.cmake) set(_test_srcs test_generation_config.cpp + test_hf_static_io_manager.cpp + test_io_manager.cpp test_model_metadata.cpp test_text_llm_runner.cpp test_text_prefiller.cpp @@ -49,32 +51,54 @@ set(_metadata_pte_files "${CMAKE_CURRENT_BINARY_DIR}/ModelMetadata_geometry_invalid_windows.pte" "${CMAKE_CURRENT_BINARY_DIR}/ModelMetadata_geometry_empty.pte" ) +# Local builds may share a host with other model exports. +set(_export_prefix) +if(EXECUTORCH_TEST_EXPORT_LOCK) + find_program(_flock flock REQUIRED) + set(_export_prefix ${_flock} ${EXECUTORCH_TEST_EXPORT_LOCK}) +endif() + add_custom_command( OUTPUT ${_metadata_pte_files} - # Run against the installed executorch (which stages exir/_serialize/*.fbs at - # install time). Do NOT prepend the source tree to PYTHONPATH: it shadows the - # installed package with a source copy that lacks the staged .fbs, breaking - # to_executorch() under a non-editable install. COMMAND - ${PYTHON_EXECUTABLE} ${CMAKE_CURRENT_SOURCE_DIR}/export_model_metadata.py + ${_export_prefix} ${PYTHON_EXECUTABLE} + ${CMAKE_CURRENT_SOURCE_DIR}/export_model_metadata.py --outdir ${CMAKE_CURRENT_BINARY_DIR} WORKING_DIRECTORY ${EXECUTORCH_ROOT} DEPENDS ${CMAKE_CURRENT_SOURCE_DIR}/export_model_metadata.py ${EXECUTORCH_ROOT}/extension/llm/export/model_metadata.py ) +add_custom_command( + OUTPUT "${CMAKE_CURRENT_BINARY_DIR}/hf_static_fixtures.stamp" + COMMAND + ${_export_prefix} ${PYTHON_EXECUTABLE} + ${CMAKE_CURRENT_SOURCE_DIR}/export_hf_static_program.py --outdir + ${CMAKE_CURRENT_BINARY_DIR} + WORKING_DIRECTORY ${EXECUTORCH_ROOT} + DEPENDS ${CMAKE_CURRENT_SOURCE_DIR}/export_hf_static_program.py +) add_custom_target( - generated_model_metadata_test_files DEPENDS ${_metadata_pte_files} + generated_model_metadata_test_files + DEPENDS ${_metadata_pte_files} + "${CMAKE_CURRENT_BINARY_DIR}/hf_static_fixtures.stamp" ) et_cxx_test( - test_runner SOURCES ${_test_srcs} EXTRA_LIBS executorch extension_llm_runner + test_runner + SOURCES + ${_test_srcs} + EXTRA_LIBS + executorch + extension_llm_runner + portable_kernels + portable_ops_lib ) add_dependencies(test_runner generated_model_metadata_test_files) set_property( TEST test_runner PROPERTY ENVIRONMENT - "ET_MODEL_METADATA_FULL_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModelMetadata_full.pte;ET_MODEL_METADATA_LAST_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModelMetadata_last.pte;ET_MODEL_METADATA_SELECTED_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModelMetadata_selected.pte;ET_MODEL_METADATA_INVALID_CONTEXT_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModelMetadata_invalid_context.pte;ET_MODEL_METADATA_INVALID_PREFILL_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModelMetadata_invalid_prefill.pte;ET_MODEL_METADATA_INVALID_VOCAB_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModelMetadata_invalid_vocab.pte;ET_MODEL_METADATA_MISSING_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModelMetadata_missing.pte;ET_MODEL_METADATA_GEOMETRY_WRONG_TYPE_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModelMetadata_geometry_wrong_type.pte;ET_MODEL_METADATA_GEOMETRY_MISMATCHED_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModelMetadata_geometry_mismatched.pte;ET_MODEL_METADATA_GEOMETRY_INVALID_HEADS_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModelMetadata_geometry_invalid_heads.pte;ET_MODEL_METADATA_GEOMETRY_INVALID_DIMS_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModelMetadata_geometry_invalid_dims.pte;ET_MODEL_METADATA_GEOMETRY_INVALID_WINDOWS_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModelMetadata_geometry_invalid_windows.pte;ET_MODEL_METADATA_GEOMETRY_EMPTY_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModelMetadata_geometry_empty.pte" + "ET_MODEL_METADATA_FULL_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModelMetadata_full.pte;ET_MODEL_METADATA_LAST_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModelMetadata_last.pte;ET_MODEL_METADATA_SELECTED_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModelMetadata_selected.pte;ET_MODEL_METADATA_INVALID_CONTEXT_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModelMetadata_invalid_context.pte;ET_MODEL_METADATA_INVALID_PREFILL_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModelMetadata_invalid_prefill.pte;ET_MODEL_METADATA_INVALID_VOCAB_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModelMetadata_invalid_vocab.pte;ET_MODEL_METADATA_MISSING_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModelMetadata_missing.pte;ET_MODEL_METADATA_GEOMETRY_WRONG_TYPE_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModelMetadata_geometry_wrong_type.pte;ET_MODEL_METADATA_GEOMETRY_MISMATCHED_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModelMetadata_geometry_mismatched.pte;ET_MODEL_METADATA_GEOMETRY_INVALID_HEADS_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModelMetadata_geometry_invalid_heads.pte;ET_MODEL_METADATA_GEOMETRY_INVALID_DIMS_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModelMetadata_geometry_invalid_dims.pte;ET_MODEL_METADATA_GEOMETRY_INVALID_WINDOWS_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModelMetadata_geometry_invalid_windows.pte;ET_MODEL_METADATA_GEOMETRY_EMPTY_PATH=${CMAKE_CURRENT_BINARY_DIR}/ModelMetadata_geometry_empty.pte;ET_HF_STATIC_TINY_PATH=${CMAKE_CURRENT_BINARY_DIR}/HfStaticTiny.pte;ET_HF_STATIC_TINY_HALF_PATH=${CMAKE_CURRENT_BINARY_DIR}/HfStaticTinyHalf.pte;ET_HF_STATIC_TINY_STATE_PATH=${CMAKE_CURRENT_BINARY_DIR}/HfStaticTinyWithState.pte;ET_HF_STATIC_FIXTURES_DIR=${CMAKE_CURRENT_BINARY_DIR}" ) # Override sanitizer to this issue: diff --git a/extension/llm/runner/test/export_hf_static_program.py b/extension/llm/runner/test/export_hf_static_program.py new file mode 100644 index 00000000000..6f77b24b692 --- /dev/null +++ b/extension/llm/runner/test/export_hf_static_program.py @@ -0,0 +1,308 @@ +# 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. + +"""Tiny float-cache decoders and malformed programs for C++ runner tests.""" + +import argparse +import os + +import torch +from executorch.exir import EdgeCompileConfig, ExecutorchBackendConfig, to_edge +from executorch.exir.passes import MemoryPlanningPass + +CONTEXT = 16 +PREFILL = 4 +VOCAB = 32 +HIDDEN = 16 +HEADS = 2 +LAYERS = 2 +PROMPT = [3, 7, 1, 9, 4, 12] +NEW_TOKENS = 6 + + +def _heads(x, weight): + return (x @ weight.t()).view(1, x.shape[1], HEADS, -1).transpose(1, 2) + + +class TinyHfStaticDecoder(torch.nn.Module): + def __init__(self, return_state=False): + super().__init__() + self.return_state = return_state + torch.manual_seed(0) + self.embed = torch.nn.Embedding(VOCAB, HIDDEN) + self.qkvo = torch.nn.ModuleList( + torch.nn.Linear(HIDDEN, 4 * HIDDEN, bias=False) for _ in range(LAYERS) + ) + self.head = torch.nn.Linear(HIDDEN, VOCAB, bias=False) + + def forward(self, input_ids, attention_mask, position_ids, *caches): + steps = input_ids.shape[1] + x = self.embed(input_ids) + for layer, keys, values, counter in zip( + self.qkvo, caches[0::3], caches[1::3], caches[2::3] + ): + slots = torch.arange(steps) + counter + q, k, v, o = layer.weight.split(HIDDEN) + keys.index_copy_(2, slots, _heads(x, k).to(keys.dtype)) + values.index_copy_(2, slots, _heads(x, v).to(values.dtype)) + counter.add_(steps) + attn = torch.nn.functional.scaled_dot_product_attention( + _heads(x, q), keys.float(), values.float(), attn_mask=attention_mask + ) + x = x + attn.transpose(1, 2).reshape(1, steps, HIDDEN) @ o.t() + logits = self.head(x) + return (logits, *caches) if self.return_state else logits + + +def _inputs(tokens, start, steps): + positions = torch.arange(start, start + steps) + mask = (torch.arange(CONTEXT)[None] <= positions[:, None])[None, None] + ids = torch.tensor([list(tokens) + [tokens[-1]] * (steps - len(tokens))]) + return ids, mask, positions[None] + + +def _caches(dtype=torch.float32): + shape = (1, HEADS, CONTEXT, HIDDEN // HEADS) + return [ + tensor + for _ in range(LAYERS) + for tensor in ( + torch.zeros(shape, dtype=dtype), + torch.zeros(shape, dtype=dtype), + torch.zeros((), dtype=torch.long), + ) + ] + + +def expected_tokens(model, dtype): + caches, start = _caches(dtype), 0 + generated = [] + with torch.no_grad(): + while start < len(PROMPT): + chunk = PROMPT[start : start + PREFILL] + for counter in caches[2::3]: + counter.fill_(start) + output = model(*_inputs(chunk, start, PREFILL), *caches) + logits = (output[0] if model.return_state else output)[0, len(chunk) - 1] + start += len(chunk) + for _ in range(NEW_TOKENS): + token = int(logits.argmax()) + generated.append(token) + for counter in caches[2::3]: + counter.fill_(start) + output = model(*_inputs([token], start, 1), *caches) + logits = (output[0] if model.return_state else output)[0, 0] + start += 1 + return generated + + +def _metadata(dtype=torch.float32): + return { + "get_io_contract": "hf_static", + "get_io_contract_version": 1, + "get_cache_dtype": "float32" if dtype == torch.float32 else "float16", + "get_n_layers": LAYERS, + "get_max_seq_len": CONTEXT, + } + + +def _lower(programs, metadata): + edge = to_edge( + programs, + compile_config=EdgeCompileConfig(_check_ir_validity=False), + constant_methods=metadata, + ) + return edge.to_executorch( + ExecutorchBackendConfig( + run_reinplace_pass=True, + memory_planning_pass=MemoryPlanningPass( + alloc_graph_input=False, alloc_graph_output=False + ), + ) + ) + + +def export(model, dtype): + programs = { + name: torch.export.export( + model, (*_inputs([0], 0, steps), *_caches(dtype)), strict=False + ) + for name, steps in (("prefill", PREFILL), ("decode", 1)) + } + metadata = _metadata(dtype) + metadata.update( + get_test_prompt=PROMPT, get_test_expected=expected_tokens(model, dtype) + ) + return _lower(programs, metadata) + + +class ContractProbe(torch.nn.Module): + """Keep the input ABI and mutable outputs without running attention.""" + + def __init__(self, first_cache, logits_shape=None, state_case=None): + super().__init__() + self.first_cache = first_cache + self.logits_shape = logits_shape + self.state_case = state_case + + def forward(self, *inputs): + for tensor in inputs[self.first_cache :]: + tensor.add_(1) + shape = self.logits_shape or (1, inputs[0].shape[1], VOCAB) + logits = torch.zeros(shape) + inputs[0].sum().float() + if self.state_case == "state_count": + return logits, inputs[self.first_cache] + if self.state_case == "state_shape": + state = list(inputs[self.first_cache :]) + state[0] = state[0].reshape(1, 4, CONTEXT, 4) + return logits, *state + return logits + + +def export_probe(case): + metadata = _metadata() + overrides = { + "newer_version": {"get_io_contract_version": 2}, + "zero_version": {"get_io_contract_version": 0}, + "version_type": {"get_io_contract_version": "1"}, + "dtype_mismatch": {"get_cache_dtype": "float16"}, + "dtype_unknown": {"get_cache_dtype": "uint16"}, + "dtype_type": {"get_cache_dtype": 1}, + "layers_mismatch": {"get_n_layers": 3}, + "layers_zero": {"get_n_layers": 0}, + "layers_type": {"get_n_layers": "2"}, + "context_zero": {"get_max_seq_len": 0}, + "context_short": {"get_max_seq_len": 8}, + "context_type": {"get_max_seq_len": "16"}, + "context_large": {"get_max_seq_len": 1 << 31}, + "context_budget_large": {"get_max_context_len": 1 << 31}, + "unknown_contract": {"get_io_contract": "hf_static_future"}, + "contract_type": {"get_io_contract": 1}, + } + metadata.update(overrides.get(case, {})) + if case.startswith("missing_"): + del metadata["get_" + case.removeprefix("missing_")] + programs = {} + for name, steps in (("prefill", PREFILL), ("decode", 1)): + inputs, first_cache, logits_shape = _probe_inputs(case, name, steps) + model = ContractProbe(first_cache, logits_shape, case) + programs[name] = torch.export.export(model, tuple(inputs), strict=False) + return _lower(programs, metadata) + + +def _probe_inputs(case, name, steps): + inputs = list(_inputs([0], 0, steps)) + _caches() + first_cache = 3 + if case in ("windows", "all_window"): + window_mask = torch.ones(1, 1, steps, 6, dtype=torch.bool) + if case == "windows": + inputs.insert(2, window_mask) + first_cache = 4 + else: + inputs[1] = window_mask + layers = LAYERS if case == "all_window" else 1 + for layer in range(layers): + for c in (first_cache + 3 * layer, first_cache + 3 * layer + 1): + inputs[c] = torch.zeros(1, HEADS, 6, HIDDEN // HEADS) + shapes = { + ("prefill", "tokens_batch"): [(0, (2, steps), torch.long)], + ("prefill", "mask_rows"): [(1, (1, 1, 1, CONTEXT), torch.bool)], + ("prefill", "mask_batch"): [(1, (2, 1, steps, CONTEXT), torch.bool)], + ("prefill", "mask_heads"): [(1, (1, 2, steps, CONTEXT), torch.bool)], + ("prefill", "mask_width"): [(1, (1, 1, steps, CONTEXT + 1), torch.bool)], + ("prefill", "position_rows"): [(2, (1, 1), torch.long)], + ("prefill", "position_batch"): [(2, (2, steps), torch.long)], + ("prefill", "kv_shape"): [ + (4, (1, HEADS, HIDDEN // HEADS, CONTEXT), torch.float32) + ], + ("decode", "method_shape"): [ + (c, (1, 4, CONTEXT, 4), torch.float32) for c in (3, 4) + ], + ("decode", "method_dtype"): [ + (c, (1, HEADS, CONTEXT, HIDDEN // HEADS * 2), torch.float16) for c in (3, 4) + ], + } + for index, shape, dtype in shapes.get((name, case), []): + inputs[index] = torch.zeros(shape, dtype=dtype) + if case == "cache_batch": + for c in (3, 4): + inputs[c] = torch.zeros(2, HEADS, CONTEXT, HIDDEN // HEADS) + logits_shape = { + ("prefill", "logits_rank"): (VOCAB,), + ("prefill", "logits_last_only"): (1, 1, VOCAB), + ("prefill", "logits_rows"): (1, steps + 1, VOCAB), + ("decode", "method_vocab"): (1, 1, VOCAB + 1), + }.get((name, case)) + return inputs, first_cache, logits_shape + + +PROBE_CASES = ( + "windows", + "all_window", + "newer_version", + "zero_version", + "version_type", + "missing_io_contract_version", + "dtype_mismatch", + "dtype_unknown", + "dtype_type", + "missing_cache_dtype", + "layers_mismatch", + "layers_zero", + "layers_type", + "missing_n_layers", + "context_zero", + "context_short", + "context_type", + "context_large", + "context_budget_large", + "missing_max_seq_len", + "unknown_contract", + "contract_type", + "tokens_batch", + "mask_rows", + "mask_batch", + "mask_heads", + "mask_width", + "position_rows", + "position_batch", + "cache_batch", + "kv_shape", + "method_shape", + "method_dtype", + "logits_rank", + "logits_last_only", + "logits_rows", + "method_vocab", + "state_count", + "state_shape", +) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--outdir", required=True) + args = parser.parse_args() + os.makedirs(args.outdir, exist_ok=True) + torch.set_num_threads(1) + torch.set_grad_enabled(False) + for name, dtype, return_state in ( + ("HfStaticTiny", torch.float32, False), + ("HfStaticTinyHalf", torch.float16, False), + ("HfStaticTinyWithState", torch.float32, True), + ): + model = TinyHfStaticDecoder(return_state).eval().requires_grad_(False) + with open(os.path.join(args.outdir, name + ".pte"), "wb") as f: + export(model, dtype).write_to_file(f) + for case in PROBE_CASES: + with open(os.path.join(args.outdir, "HfStatic_" + case + ".pte"), "wb") as f: + export_probe(case).write_to_file(f) + with open(os.path.join(args.outdir, "hf_static_fixtures.stamp"), "w") as f: + f.write("float-cache v1 fixtures\n") + + +if __name__ == "__main__": + main() diff --git a/extension/llm/runner/test/targets.bzl b/extension/llm/runner/test/targets.bzl index 08044de2d35..76523332e8b 100644 --- a/extension/llm/runner/test/targets.bzl +++ b/extension/llm/runner/test/targets.bzl @@ -27,6 +27,14 @@ def define_common_targets(): ], ) + runtime.cxx_test( + name = "test_io_manager", + srcs = ["test_io_manager.cpp"], + deps = [ + "//executorch/extension/llm/runner/io_manager:io_manager", + ], + ) + runtime.cxx_test( name = "test_text_prefiller", srcs = ["test_text_prefiller.cpp"], diff --git a/extension/llm/runner/test/test_hf_static_io_manager.cpp b/extension/llm/runner/test/test_hf_static_io_manager.cpp new file mode 100644 index 00000000000..9c69a012bdd --- /dev/null +++ b/extension/llm/runner/test/test_hf_static_io_manager.cpp @@ -0,0 +1,590 @@ +/* + * 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. + */ + +// HfStaticIOManager drives a program with Hugging Face's static-cache ABI +// (export_hf_static_program.py): it owns the caches, which the program +// writes in place, and builds each step's tokens, masks and positions. The +// program's greedy tokens can only match the expected ones if every step's +// writes land in the IOManager's buffers. + +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include + +using executorch::extension::from_blob; +using executorch::extension::Module; +using executorch::extension::llm::HfStaticIOManager; +using executorch::runtime::Error; + +namespace { + +constexpr int64_t kPrefill = 4; + +std::vector int_list(Module& module, const char* method) { + auto result = module.execute(method); + EXPECT_EQ(result.error(), Error::Ok); + std::vector values; + if (!result.ok()) { + return {}; + } + for (const auto& value : *result) { + values.push_back(value.toInt()); + } + return values; +} + +// Exercise both float32 and float16 cache storage. +class HfStaticIOManagerTest : public ::testing::TestWithParam { + protected: + void SetUp() override { + const char* path = std::getenv(GetParam()); + ASSERT_NE(path, nullptr); + module_ = std::make_unique(path); + io_ = std::make_unique(*module_); + auto methods = module_->method_names(); + ASSERT_EQ(methods.error(), Error::Ok); + std::tie(prefill_, decode_) = io_->method_names(*methods); + ASSERT_EQ(io_->load(prefill_, decode_), Error::Ok); + } + + // Runs `tokens` at `start` through `method`; returns the greedy next token. + int64_t + step(const std::string& method, std::vector tokens, int64_t start) { + auto ids = from_blob( + tokens.data(), + {1, static_cast(tokens.size())}, + executorch::aten::ScalarType::Long); + auto pos = from_blob(&start, {1}, executorch::aten::ScalarType::Long); + auto inputs = method == prefill_ ? io_->prepare_prefill(ids, pos, method) + : io_->prepare_decode(ids, pos, method); + EXPECT_EQ(inputs.error(), Error::Ok); + if (!inputs.ok()) { + return -1; + } + auto outputs = module_->execute(method, *inputs); + EXPECT_EQ(outputs.error(), Error::Ok); + if (!outputs.ok()) { + return -1; + } + if (outputs->size() == 13) { + for (size_t c = 0; c < 6; ++c) { + EXPECT_EQ( + (*outputs)[7 + c].toTensor().const_data_ptr(), + (*inputs)[3 + c].toTensor().const_data_ptr()); + } + } + EXPECT_EQ(io_->update_decode(*outputs, method), Error::Ok); + auto logits = io_->logits(*outputs, method); + EXPECT_EQ(logits.error(), Error::Ok); + if (!logits.ok()) { + return -1; + } + const float* row = logits->const_data_ptr(); + return std::max_element(row, row + logits->numel()) - row; + } + + std::vector generate(const std::string& prompt_method) { + const auto prompt = int_list(*module_, "get_test_prompt"); + const auto expected = int_list(*module_, "get_test_expected"); + const int64_t chunk = io_->max_tokens_per_step(prompt_method); + int64_t start = 0, next = 0; + while (start < static_cast(prompt.size())) { + const auto end = std::min(prompt.size(), start + chunk); + std::vector part(prompt.begin() + start, prompt.begin() + end); + next = step(prompt_method, part, start); + start = end; + } + std::vector generated; + for (size_t i = 0; i < expected.size(); ++i) { + generated.push_back(next); + next = step(decode_, {next}, start++); + } + return generated; + } + + std::unique_ptr module_; + std::unique_ptr io_; + std::string prefill_, decode_; +}; + +TEST_P(HfStaticIOManagerTest, PicksPrefillAndDecode) { + EXPECT_EQ(prefill_, "prefill"); + EXPECT_EQ(decode_, "decode"); + EXPECT_EQ(io_->max_tokens_per_step(prefill_), kPrefill); + EXPECT_EQ(io_->max_tokens_per_step(decode_), 1); +} + +TEST_P(HfStaticIOManagerTest, GeneratesExpectedTokens) { + EXPECT_EQ(generate(prefill_), int_list(*module_, "get_test_expected")); +} + +TEST_P(HfStaticIOManagerTest, DecodeOnlyPromptMatches) { + // The same caches serve both methods: a prompt fed one token at a time + // through decode gives the same continuation. + EXPECT_EQ(generate(decode_), int_list(*module_, "get_test_expected")); +} + +TEST_P(HfStaticIOManagerTest, ResetClearsTheCaches) { + const auto first = generate(prefill_); + ASSERT_EQ(io_->reset(prefill_, decode_), Error::Ok); + EXPECT_EQ(generate(prefill_), first); +} + +TEST_P(HfStaticIOManagerTest, RejectsTooManyTokens) { + std::vector tokens(kPrefill + 1, 1); + int64_t start = 0; + auto ids = from_blob( + tokens.data(), + {1, static_cast(tokens.size())}, + executorch::aten::ScalarType::Long); + auto pos = from_blob(&start, {1}, executorch::aten::ScalarType::Long); + EXPECT_EQ( + io_->prepare_prefill(ids, pos, prefill_).error(), Error::InvalidArgument); +} + +TEST_P(HfStaticIOManagerTest, RejectsAPromptChunkPastTheCache) { + // A prompt chunk writes all its rows, padding included. + std::vector tokens = {1}; + int64_t start = 16 - kPrefill + 1; + auto ids = + from_blob(tokens.data(), {1, 1}, executorch::aten::ScalarType::Long); + auto pos = from_blob(&start, {1}, executorch::aten::ScalarType::Long); + EXPECT_EQ( + io_->prepare_prefill(ids, pos, prefill_).error(), Error::OutOfResources); +} + +TEST_P(HfStaticIOManagerTest, RejectsNegativePositions) { + std::vector tokens = {1}; + int64_t start = -1; + auto ids = + from_blob(tokens.data(), {1, 1}, executorch::aten::ScalarType::Long); + auto pos = from_blob(&start, {1}, executorch::aten::ScalarType::Long); + EXPECT_EQ( + io_->prepare_decode(ids, pos, decode_).error(), Error::InvalidArgument); +} + +TEST_P(HfStaticIOManagerTest, RejectsPositionsPastTheCache) { + std::vector tokens = {1}; + int64_t start = 16; // the program's context length + auto ids = + from_blob(tokens.data(), {1, 1}, executorch::aten::ScalarType::Long); + auto pos = from_blob(&start, {1}, executorch::aten::ScalarType::Long); + EXPECT_EQ( + io_->prepare_decode(ids, pos, decode_).error(), Error::OutOfResources); +} + +INSTANTIATE_TEST_SUITE_P( + Caches, + HfStaticIOManagerTest, + ::testing::Values( + "ET_HF_STATIC_TINY_PATH", + "ET_HF_STATIC_TINY_HALF_PATH", + "ET_HF_STATIC_TINY_STATE_PATH")); + +TEST_P(HfStaticIOManagerTest, PrepareSetsEveryCounterFromStart) { + int64_t token = 1, start = 7; + auto ids = from_blob(&token, {1, 1}, executorch::aten::ScalarType::Long); + auto pos = from_blob(&start, {1}, executorch::aten::ScalarType::Long); + auto inputs = io_->prepare_prefill(ids, pos, prefill_); + ASSERT_TRUE(inputs.ok()); + for (size_t c = 5; c < inputs->size(); c += 3) { + EXPECT_EQ((*inputs)[c].toTensor().const_data_ptr()[0], start); + (*inputs)[c].toTensor().mutable_data_ptr()[0] = 100; + } + start = 9; + auto decode_inputs = io_->prepare_decode(ids, pos, decode_); + ASSERT_TRUE(decode_inputs.ok()); + for (size_t c = 5; c < decode_inputs->size(); c += 3) { + EXPECT_EQ( + (*decode_inputs)[c].toTensor().const_data_ptr()[0], start); + } +} + +TEST_P(HfStaticIOManagerTest, GeneratesAgainWithoutAnExplicitCacheReset) { + const auto first = generate(prefill_); + EXPECT_EQ(first, int_list(*module_, "get_test_expected")); + EXPECT_EQ(generate(prefill_), first); +} + +std::string fixture_path(const std::string& name) { + const auto* dir = std::getenv("ET_HF_STATIC_FIXTURES_DIR"); + return std::string(dir ? dir : "") + "/" + name + ".pte"; +} + +class HfStaticMalformedTest : public ::testing::TestWithParam {}; + +TEST_P(HfStaticMalformedTest, RejectsTheProgramAtLoad) { + Module module(fixture_path(std::string("HfStatic_") + GetParam())); + HfStaticIOManager io(module); + EXPECT_EQ(io.load("prefill", "decode"), Error::InvalidArgument); +} + +INSTANTIATE_TEST_SUITE_P( + Contract, + HfStaticMalformedTest, + ::testing::Values( + "zero_version", + "version_type", + "missing_io_contract_version", + "dtype_mismatch", + "dtype_unknown", + "dtype_type", + "missing_cache_dtype", + "layers_mismatch", + "layers_zero", + "layers_type", + "missing_n_layers", + "context_zero", + "context_short", + "context_type", + "context_large", + "missing_max_seq_len", + "tokens_batch", + "mask_rows", + "mask_batch", + "mask_heads", + "mask_width", + "position_rows", + "position_batch", + "cache_batch", + "kv_shape", + "method_shape", + "method_dtype", + "logits_rank", + "logits_last_only", + "logits_rows", + "method_vocab", + "state_count", + "state_shape"), + [](const ::testing::TestParamInfo& info) { + return info.param; + }); + +TEST(HfStaticContractTest, RejectsNewerContractVersion) { + Module module(fixture_path("HfStatic_newer_version")); + HfStaticIOManager io(module); + EXPECT_EQ(io.load("prefill", "decode"), Error::NotSupported); +} + +TEST(HfStaticContractTest, BuildsMultipleMasksAndMarksOnlyWindowPadding) { + Module module(fixture_path("HfStatic_windows")); + HfStaticIOManager io(module); + ASSERT_EQ(io.load("prefill", "decode"), Error::Ok); + std::vector tokens = {7, 9}; + int64_t start = 5; + auto ids = + from_blob(tokens.data(), {1, 2}, executorch::aten::ScalarType::Long); + auto pos = from_blob(&start, {1}, executorch::aten::ScalarType::Long); + auto inputs = io.prepare_prefill(ids, pos, "prefill"); + ASSERT_TRUE(inputs.ok()); + EXPECT_EQ((*inputs)[0].toTensor().const_data_ptr()[3], 9); + EXPECT_EQ((*inputs)[3].toTensor().const_data_ptr()[3], 8); + for (size_t m : {1, 2}) { + const auto& mask = (*inputs)[m].toTensor(); + const int64_t width = mask.size(3); + for (int64_t r = 0; r < kPrefill; ++r) { + for (int64_t c = 0; c < width; ++c) { + const bool visible = !(m == 2 && r >= 2) && c <= start + r; + EXPECT_EQ(mask.const_data_ptr()[r * width + c], visible); + } + } + } + auto outputs = module.execute("prefill", *inputs); + ASSERT_TRUE(outputs.ok()); + for (size_t c = 4; c < inputs->size(); ++c) { + EXPECT_EQ( + (*outputs)[c - 4].toTensor().const_data_ptr(), + (*inputs)[c].toTensor().const_data_ptr()); + } +} + +TEST(HfStaticContractTest, AllWindowCachesUseMetadataContext) { + Module module(fixture_path("HfStatic_all_window")); + HfStaticIOManager io(module); + ASSERT_EQ(io.load("prefill", "decode"), Error::Ok); + int64_t token = 1, start = 10; + auto ids = from_blob(&token, {1, 1}, executorch::aten::ScalarType::Long); + auto pos = from_blob(&start, {1}, executorch::aten::ScalarType::Long); + auto inputs = io.prepare_prefill(ids, pos, "prefill"); + ASSERT_TRUE(inputs.ok()); + const auto& mask = (*inputs)[1].toTensor(); + for (int64_t r = 0; r < kPrefill; ++r) { + for (int64_t c = 0; c < mask.size(3); ++c) { + EXPECT_EQ(mask.const_data_ptr()[r * mask.size(3) + c], r == 0); + } + } +} + +class TinyTokenizer : public ::tokenizers::Tokenizer { + public: + TinyTokenizer() { + initialized_ = true; + vocab_size_ = 32; + eos_tok_ = 100; + } + ::tokenizers::Error load(const std::string&) override { + return ::tokenizers::Error::Ok; + } + ::tokenizers::Result> + encode(const std::string& text, int8_t, int8_t) const override { + if (text == "tail") { + return std::vector{7, 9}; + } + return std::vector{3, 7, 1, 9, 4, 12}; + } + ::tokenizers::Result decode(uint64_t, uint64_t token, bool) + const override { + return std::to_string(token); + } + ::tokenizers::Result id_to_piece(uint64_t token) const override { + return std::to_string(token); + } + ::tokenizers::Result piece_to_id( + const std::string&) const override { + return uint64_t{0}; + } +}; + +class ObservedHfStaticIOManager : public HfStaticIOManager { + public: + explicit ObservedHfStaticIOManager(Module& module) + : HfStaticIOManager(module) {} + Error reset(const std::string& prefill, const std::string& decode) override { + ++resets; + return HfStaticIOManager::reset(prefill, decode); + } + int resets = 0; +}; + +TEST_P(HfStaticIOManagerTest, RunnerResetRegeneratesIdenticalTokens) { + const char* path = std::getenv(GetParam()); + auto module = std::make_unique(path); + auto io = std::make_unique(*module); + auto* observed = io.get(); + ASSERT_EQ(io->load("prefill", "decode"), Error::Ok); + auto tokenizer = std::make_unique(); + auto metadata = executorch::extension::llm::get_llm_metadata( + tokenizer.get(), module.get()); + ASSERT_TRUE(metadata.ok()); + auto stats = std::make_unique(); + auto prefill = + std::make_unique( + module.get(), io.get(), "prefill"); + auto decode = std::make_unique( + module.get(), io.get(), "decode"); + auto prefiller = std::make_unique( + prefill.get(), true, true, kPrefill); + auto generator = + std::make_unique( + tokenizer.get(), + decode.get(), + true, + std::make_unique>(), + stats.get()); + executorch::extension::llm::TextLLMRunner runner( + *metadata, + std::move(tokenizer), + std::move(module), + std::move(prefill), + std::move(prefiller), + std::move(io), + std::move(generator), + std::move(stats), + -1.0f, + std::move(decode)); + executorch::extension::llm::GenerationConfig config; + config.echo = false; + config.ignore_eos = true; + config.max_new_tokens = 6; + config.temperature = 0.0f; + std::vector first, second; + ASSERT_EQ( + runner.generate( + "prompt", + config, + [&](const std::string& s) { first.push_back(std::stoll(s)); }), + Error::Ok); + runner.reset(); + EXPECT_EQ(observed->resets, 1); + ASSERT_EQ( + runner.generate( + "prompt", + config, + [&](const std::string& s) { second.push_back(std::stoll(s)); }), + Error::Ok); + EXPECT_EQ(first.size(), 6); + EXPECT_EQ(first, int_list(*module_, "get_test_expected")); + EXPECT_EQ(second, first); +} + +TEST(HfStaticRunnerTest, LoadUsesTheSelectedMethodNames) { + auto module = std::make_unique(fixture_path("HfStaticTiny")); + auto io = std::make_unique(*module); + auto tokenizer = std::make_unique(); + auto stats = std::make_unique(); + auto prefill = + std::make_unique( + module.get(), io.get(), "prefill"); + auto decode = std::make_unique( + module.get(), io.get(), "decode"); + auto prefiller = std::make_unique( + prefill.get(), true, true, kPrefill); + auto generator = + std::make_unique( + tokenizer.get(), + decode.get(), + true, + std::make_unique>(), + stats.get()); + executorch::extension::llm::TextLLMRunner runner( + {}, + std::move(tokenizer), + std::move(module), + std::move(prefill), + std::move(prefiller), + std::move(io), + std::move(generator), + std::move(stats), + -1.0f, + std::move(decode)); + EXPECT_EQ(runner.load(), Error::Ok); +} + +TEST(HfStaticRunnerTest, RejectsUnknownAndMalformedContractLabels) { + for (const auto* name : {"unknown_contract", "contract_type"}) { + auto runner = executorch::extension::llm::create_text_llm_runner( + fixture_path(std::string("HfStatic_") + name), + std::make_unique()); + EXPECT_EQ(runner, nullptr); + } +} + +TEST(HfStaticRunnerTest, RejectsMalformedRuntimeMetadataWithoutAborting) { + for (const auto* name : + {"context_type", + "layers_type", + "version_type", + "dtype_type", + "newer_version"}) { + auto runner = executorch::extension::llm::create_text_llm_runner( + fixture_path(std::string("HfStatic_") + name), + std::make_unique()); + EXPECT_EQ(runner, nullptr); + } +} + +TEST(HfStaticRunnerTest, RejectsOversizedExplicitContextBudget) { + auto runner = executorch::extension::llm::create_text_llm_runner( + fixture_path("HfStatic_context_budget_large"), + std::make_unique()); + EXPECT_EQ(runner, nullptr); +} + +TEST_P(HfStaticIOManagerTest, RunnerUsesDecodeForNearFullContextTail) { + const char* path = std::getenv(GetParam()); + auto runner = executorch::extension::llm::create_text_llm_runner( + path, std::make_unique()); + ASSERT_NE(runner, nullptr); + executorch::extension::llm::GenerationConfig config; + config.echo = false; + config.ignore_eos = true; + config.temperature = 0.0f; + config.max_new_tokens = 8; + ASSERT_EQ(runner->generate("prompt", config), Error::Ok); // position 13 + config.max_new_tokens = 1; + std::vector tail; + EXPECT_EQ( + runner->generate( + "tail", + config, + [&](const std::string& s) { tail.push_back(std::stoll(s)); }), + Error::Ok); + EXPECT_EQ(tail.size(), 1); + // Independent decode-only prefill, with identical history. + ASSERT_EQ(io_->reset(prefill_, decode_), Error::Ok); + const auto prompt = int_list(*module_, "get_test_prompt"); + int64_t next = 0, start = 0; + for (const auto token : prompt) { + next = step(decode_, {token}, start++); + } + for (int i = 0; i < 7; ++i) { + next = step(decode_, {next}, start++); + } + next = step(decode_, {7}, start++); + next = step(decode_, {9}, start++); + EXPECT_EQ(tail, std::vector{next}); +} + +struct AllocatorCounts { + size_t allocated = 0, bound = 0, freed = 0; + bool fallback = false; +}; + +class CountingAllocator : public executorch::extension::llm::IOBufferAllocator { + public: + explicit CountingAllocator(AllocatorCounts& counts) : counts_(counts) {} + void* allocate(size_t nbytes, size_t alignment) override { + ++counts_.allocated; + if (counts_.fallback) { + return nullptr; + } + void* block = ::operator new(nbytes, std::align_val_t(alignment)); + blocks_[block] = nbytes; + return block; + } + void bind(void* block, void* tensor_data) override { + ASSERT_TRUE(blocks_.count(block)); + auto offset = + static_cast(tensor_data) - static_cast(block); + EXPECT_GE(offset, 0); + EXPECT_LT(static_cast(offset), blocks_.at(block)); + EXPECT_EQ(reinterpret_cast(tensor_data) % 64, 0); + ++counts_.bound; + } + void deallocate(void* block) override { + EXPECT_EQ(blocks_.erase(block), 1); + ++counts_.freed; + ::operator delete(block, std::align_val_t(64)); + } + + private: + AllocatorCounts& counts_; + std::unordered_map blocks_; +}; + +TEST(HfStaticContractTest, AllocatorBindsAndFreesOnlyItsOwnBlocks) { + auto previous = executorch::extension::llm::io_buffer_allocator_factory(); + for (bool fallback : {false, true}) { + AllocatorCounts counts; + counts.fallback = fallback; + executorch::extension::llm::set_io_buffer_allocator( + [&counts] { return std::make_unique(counts); }); + { + Module module(fixture_path("HfStaticTiny")); + HfStaticIOManager io(module); + EXPECT_EQ(io.load("prefill", "decode"), Error::Ok); + EXPECT_EQ(counts.allocated, 3); + EXPECT_EQ(counts.bound, fallback ? 0 : 14); + } + EXPECT_EQ(counts.freed, fallback ? 0 : counts.allocated); + } + executorch::extension::llm::set_io_buffer_allocator(std::move(previous)); +} + +} // namespace diff --git a/extension/llm/runner/test/test_io_manager.cpp b/extension/llm/runner/test/test_io_manager.cpp new file mode 100644 index 00000000000..ee6b2bf23b4 --- /dev/null +++ b/extension/llm/runner/test/test_io_manager.cpp @@ -0,0 +1,57 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#include +#include + +using executorch::extension::Module; +using executorch::extension::llm::io_manager_registry; +using executorch::extension::llm::IOManager; +using executorch::extension::llm::make_io_manager; +using executorch::extension::llm::register_io_manager; +using executorch::runtime::Error; + +TEST(IOManagerRegistryTest, NoLabelUsesTheDefaultManager) { + Module module(""); + auto result = make_io_manager(module, ""); + ASSERT_TRUE(result.ok()); + EXPECT_NE(*result, nullptr); +} + +TEST(IOManagerRegistryTest, UnknownLabelReturnsAnError) { + Module module(""); + EXPECT_EQ( + make_io_manager(module, "unregistered_test_contract").error(), + Error::NotSupported); +} + +TEST(IOManagerRegistryTest, RegisteredFactoryIsUsedAndCannotBeReplaced) { + Module module(""); + int calls = 0; + const std::string label = "registry_test_contract"; + EXPECT_TRUE(register_io_manager(label, [&](Module& bound) { + EXPECT_EQ(&bound, &module); + ++calls; + return std::make_unique(bound); + })); + EXPECT_FALSE(register_io_manager( + label, [](Module&) { return std::unique_ptr{}; })); + auto result = make_io_manager(module, label); + EXPECT_TRUE(result.ok()); + EXPECT_EQ(calls, 1); + io_manager_registry().erase(label); +} + +TEST(IOManagerRegistryTest, NullFactoryResultReturnsAnError) { + Module module(""); + const std::string label = "null_test_contract"; + ASSERT_TRUE(register_io_manager( + label, [](Module&) { return std::unique_ptr{}; })); + EXPECT_EQ(make_io_manager(module, label).error(), Error::InvalidState); + io_manager_registry().erase(label); +} diff --git a/extension/llm/runner/test/test_text_llm_runner.cpp b/extension/llm/runner/test/test_text_llm_runner.cpp index f6f5137a90f..c64433515c6 100644 --- a/extension/llm/runner/test/test_text_llm_runner.cpp +++ b/extension/llm/runner/test/test_text_llm_runner.cpp @@ -78,7 +78,8 @@ class MockModule : public ::executorch::extension::Module { class MockTextDecoderRunner : public TextDecoderRunner { public: - MockTextDecoderRunner() : TextDecoderRunner(nullptr, nullptr) {} + explicit MockTextDecoderRunner(std::string method = "forward") + : TextDecoderRunner(nullptr, nullptr, std::move(method)) {} MOCK_METHOD( Result, step, @@ -761,4 +762,52 @@ TEST_F(RunnerTest, TextTokenGeneratorWithoutProcessorPicksArgmax) { EXPECT_EQ(generated_tokens, expected); } +class MockNamedIOManager : public executorch::extension::llm::IOManager { + public: + explicit MockNamedIOManager(executorch::extension::Module& module) + : IOManager(module) {} + MOCK_METHOD( + Error, + load, + (const std::string&, const std::string&), + (override)); + MOCK_METHOD( + Error, + reset, + (const std::string&, const std::string&), + (override)); +}; + +TEST_F(RunnerTest, LoadAndResetUseSelectedIOManagerMethods) { + auto module = std::make_unique(); + auto io = std::make_unique(*module); + EXPECT_CALL(*io, load("prompt_method", "decode_method")) + .WillOnce(Return(Error::Ok)); + EXPECT_CALL(*io, reset("prompt_method", "decode_method")) + .WillOnce(Return(Error::Ok)); + auto tokenizer = createMockTokenizer(); + auto prefill = std::make_unique("prompt_method"); + auto decode = std::make_unique("decode_method"); + ON_CALL(*prefill, is_method_loaded()).WillByDefault(Return(false)); + ON_CALL(*prefill, load()).WillByDefault(Return(Error::Ok)); + ON_CALL(*decode, load()).WillByDefault(Return(Error::Ok)); + auto prefiller = createMockTextPrefiller(prefill.get()); + auto stats = std::make_unique(); + auto generator = + createTextTokenGenerator(tokenizer.get(), decode.get(), stats.get()); + TextLLMRunner runner( + createDefaultMetadata(), + std::move(tokenizer), + std::move(module), + std::move(prefill), + std::move(prefiller), + std::move(io), + std::move(generator), + std::move(stats), + -1.0f, + std::move(decode)); + EXPECT_EQ(runner.load(), Error::Ok); + runner.reset(); +} + } // namespace diff --git a/extension/llm/runner/test/test_text_prefiller.cpp b/extension/llm/runner/test/test_text_prefiller.cpp index e859db36346..a44bf79d9fd 100644 --- a/extension/llm/runner/test/test_text_prefiller.cpp +++ b/extension/llm/runner/test/test_text_prefiller.cpp @@ -31,6 +31,7 @@ class MockTextDecoderRunner : public TextDecoderRunner { step, (executorch::extension::TensorPtr&, int64_t), ()); + MOCK_METHOD(int64_t, max_tokens_per_step, (int64_t), (const, override)); MOCK_METHOD(bool, is_method_loaded, (), ()); MOCK_METHOD(Result, prefill, (std::vector&, int64_t), ()); MOCK_METHOD(::executorch::runtime::Error, load, (), ()); @@ -42,6 +43,8 @@ class TextPrefillerTest : public Test { void SetUp() override { executorch::runtime::runtime_init(); // Set up default behavior for the text decoder runner + ON_CALL(text_decoder_runner_, max_tokens_per_step(_)) + .WillByDefault(Return(-1)); ON_CALL(text_decoder_runner_, is_method_loaded()) .WillByDefault(Return(true)); ON_CALL(text_decoder_runner_, step) @@ -515,4 +518,58 @@ TEST_F( // initial_pos + num_tokens, not double-incremented EXPECT_EQ(start_pos, 108); // 100 + 8 tokens } +TEST_F(TextPrefillerTest, FallsBackToDecodeForAPaddedTail) { + MockTextDecoderRunner decode; + TextPrefiller prefiller(&text_decoder_runner_, true, true, 4, &decode); + std::vector tokens = {7, 9}; + int64_t start = 13; + EXPECT_CALL(text_decoder_runner_, max_tokens_per_step(13)) + .WillOnce(Return(0)); + EXPECT_CALL(text_decoder_runner_, step(_, _)).Times(0); + EXPECT_CALL(decode, load()).WillOnce(Return(Error::Ok)); + InSequence sequence; + EXPECT_CALL(decode, step(_, 13)).WillOnce([&](auto& input, int64_t) { + EXPECT_EQ(input->numel(), 1); + EXPECT_EQ(input->template const_data_ptr()[0], 7); + return Result(tensor); + }); + EXPECT_CALL(decode, step(_, 14)).WillOnce([&](auto& input, int64_t) { + EXPECT_EQ(input->template const_data_ptr()[0], 9); + return Result(tensor); + }); + auto result = prefiller.prefill(tokens, start); + ASSERT_TRUE(result.ok()); + EXPECT_EQ(*result, 3); + EXPECT_EQ(start, 15); +} + +TEST_F(TextPrefillerTest, DoesNotRetryOtherErrorsAsDecode) { + MockTextDecoderRunner decode; + TextPrefiller prefiller(&text_decoder_runner_, true, true, 4, &decode); + std::vector tokens = {7, 9}; + int64_t start = 13; + EXPECT_CALL(text_decoder_runner_, step(_, 13)) + .WillOnce([](auto&, int64_t) -> Result { + return Error::InvalidArgument; + }); + EXPECT_CALL(decode, step(_, _)).Times(0); + EXPECT_EQ(prefiller.prefill(tokens, start).error(), Error::InvalidArgument); + EXPECT_EQ(start, 13); +} + +TEST_F(TextPrefillerTest, DoesNotRetryExecutionResourceErrors) { + MockTextDecoderRunner decode; + TextPrefiller prefiller(&text_decoder_runner_, true, true, 4, &decode); + std::vector tokens = {7, 9}; + int64_t start = 13; + EXPECT_CALL(text_decoder_runner_, step(_, 13)) + .WillOnce([](auto&, int64_t) -> Result { + return Error::OutOfResources; + }); + EXPECT_CALL(decode, load()).Times(0); + EXPECT_CALL(decode, step(_, _)).Times(0); + EXPECT_EQ(prefiller.prefill(tokens, start).error(), Error::OutOfResources); + EXPECT_EQ(start, 13); +} + } // namespace diff --git a/extension/llm/runner/text_decoder_runner.cpp b/extension/llm/runner/text_decoder_runner.cpp index bb64d79dd2e..d533a35dd96 100644 --- a/extension/llm/runner/text_decoder_runner.cpp +++ b/extension/llm/runner/text_decoder_runner.cpp @@ -81,15 +81,9 @@ ::executorch::runtime::Result TextDecoderRunner::step( io_manager_->update_decode(outputs_res.get(), method_name_); ET_CHECK_OK_OR_RETURN_ERROR(update_err); - ET_CHECK_MSG( - outputs_res.get().size() == 1, - "More than one output returned from executing LLM."); - ET_CHECK_MSG( - outputs_res.get()[0].isTensor(), - "Non Tensor Output returned from executing LLM"); - - // Return the logits tensor - return outputs_res.get()[0].toTensor(); + // The IOManager knows which output holds the logits (a model may also + // return its updated cache). + return io_manager_->logits(outputs_res.get(), method_name_); } else { // no kv cache (void)start_pos; // unused diff --git a/extension/llm/runner/text_decoder_runner.h b/extension/llm/runner/text_decoder_runner.h index 4dbce70a830..99e7e70a820 100644 --- a/extension/llm/runner/text_decoder_runner.h +++ b/extension/llm/runner/text_decoder_runner.h @@ -65,6 +65,13 @@ class ET_EXPERIMENTAL TextDecoderRunner { return module_->is_method_loaded(method_name_); } + /// Side-effect-free capacity for a step at start_pos; -1 means unlimited. + virtual int64_t max_tokens_per_step(int64_t start_pos) const { + return io_manager_ + ? io_manager_->max_tokens_per_step(method_name_, start_pos) + : -1; + } + /** * Get the method name used by this runner. * @return The method name. diff --git a/extension/llm/runner/text_llm_runner.cpp b/extension/llm/runner/text_llm_runner.cpp index ce4f6a08ea2..f6185afaa4c 100644 --- a/extension/llm/runner/text_llm_runner.cpp +++ b/extension/llm/runner/text_llm_runner.cpp @@ -39,11 +39,11 @@ TextLLMRunner::TextLLMRunner( std::unique_ptr decode_text_decoder_runner) : tokenizer_(std::move(tokenizer)), metadata_(std::move(metadata)), + io_manager_(std::move(io_manager)), module_(std::move(module)), text_decoder_runner_(std::move(text_decoder_runner)), decode_text_decoder_runner_(std::move(decode_text_decoder_runner)), text_prefiller_(std::move(text_prefiller)), - io_manager_(std::move(io_manager)), text_token_generator_(std::move(text_token_generator)), stats_(std::move(stats)), temperature_(temperature), @@ -58,7 +58,10 @@ Error TextLLMRunner::load() { return Error::Ok; } ET_CHECK_OK_OR_RETURN_ERROR(text_prefiller_->load()); - ET_CHECK_OK_OR_RETURN_ERROR(io_manager_->load()); + ET_CHECK_OK_OR_RETURN_ERROR(io_manager_->load( + text_decoder_runner_->method_name(), + decode_text_decoder_runner_ ? decode_text_decoder_runner_->method_name() + : text_decoder_runner_->method_name())); ET_CHECK_OK_OR_RETURN_ERROR(text_token_generator_->load()); return Error::Ok; } @@ -352,6 +355,17 @@ void TextLLMRunner::stop() { } void TextLLMRunner::reset() { + const auto error = io_manager_->reset( + text_decoder_runner_->method_name(), + decode_text_decoder_runner_ ? decode_text_decoder_runner_->method_name() + : text_decoder_runner_->method_name()); + if (error != Error::Ok) { + ET_LOG( + Error, + "Failed to reset the IOManager: 0x%x.", + static_cast(error)); + return; + } stats_->reset(); pos_ = 0; prefill_next_token_.reset(); diff --git a/extension/llm/runner/text_llm_runner.h b/extension/llm/runner/text_llm_runner.h index 195348a7a3d..f0ff48553d7 100644 --- a/extension/llm/runner/text_llm_runner.h +++ b/extension/llm/runner/text_llm_runner.h @@ -167,6 +167,9 @@ class ET_EXPERIMENTAL TextLLMRunner : public IRunner { // Components std::unique_ptr<::tokenizers::Tokenizer> tokenizer_; std::unordered_map metadata_; + // Outlives module_: a delegate may hold the I/O buffers it allocated (e.g. + // registered shared memory) until the module is destroyed. + std::unique_ptr io_manager_; std::unique_ptr<::executorch::extension::Module> module_; // Manage module's lifecycle, make sure it outlives // text_decoder_runner_. @@ -179,7 +182,6 @@ class ET_EXPERIMENTAL TextLLMRunner : public IRunner { // decode method, in which case // text_decoder_runner_ drives prefill only. std::unique_ptr text_prefiller_; - std::unique_ptr io_manager_; std::unique_ptr text_token_generator_; // Stats diff --git a/extension/llm/runner/text_prefiller.cpp b/extension/llm/runner/text_prefiller.cpp index a3bfcced4c7..9aa37f01575 100644 --- a/extension/llm/runner/text_prefiller.cpp +++ b/extension/llm/runner/text_prefiller.cpp @@ -20,8 +20,10 @@ TextPrefiller::TextPrefiller( TextDecoderRunner* text_decoder_runner, bool use_kv_cache, bool enable_parallel_prefill, - int64_t max_seq_len) + int64_t max_seq_len, + TextDecoderRunner* decode_runner) : text_decoder_runner_(text_decoder_runner), + decode_runner_(decode_runner), use_kv_cache_(use_kv_cache), enable_parallel_prefill_(enable_parallel_prefill), max_seq_len_(max_seq_len > 0 ? max_seq_len : 128) {} @@ -106,61 +108,39 @@ ::executorch::runtime::Result TextPrefiller::prefill_chunk( // When kv cache is not used, start pos is ignored int32_t num_prompt_tokens = prompt_tokens.size(); - // store the token - uint64_t cur_token; if (enable_parallel_prefill_ || !use_kv_cache_) { - // initialize tensor wrappers - auto tokens = from_blob( - prompt_tokens.data(), - {1, num_prompt_tokens}, - executorch::aten::ScalarType::Long); - - auto outputs_res = text_decoder_runner_->step(tokens, start_pos); - - ET_CHECK_OK_OR_RETURN_ERROR(outputs_res.error()); - ET_LOG( - Info, "Prefill token result numel(): %zu", outputs_res.get().numel()); - - start_pos += num_prompt_tokens; - cur_token = - text_decoder_runner_->logits_to_token(outputs_res.get(), temperature); - } else { // sequential prefill - int64_t pos = 0; // position in the sequence - // NOLINTNEXTLINE(facebook-hte-ParameterUncheckedArrayBounds) - cur_token = prompt_tokens[0]; - - // initialize tensor wrappers - auto tokens = - from_blob(&cur_token, {1, 1}, executorch::aten::ScalarType::Long); - - // run the first token and get back logits tensor. Assuming the first token - // is bos so don't callback. - auto logits_result = text_decoder_runner_->step(tokens, start_pos); - if (!logits_result.ok()) { - return logits_result.error(); + const auto capacity = text_decoder_runner_->max_tokens_per_step(start_pos); + if (capacity != 0) { + auto tokens = from_blob( + prompt_tokens.data(), + {1, num_prompt_tokens}, + executorch::aten::ScalarType::Long); + auto outputs = text_decoder_runner_->step(tokens, start_pos); + ET_CHECK_OK_OR_RETURN_ERROR(outputs.error()); + start_pos += num_prompt_tokens; + return text_decoder_runner_->logits_to_token(*outputs, temperature); } - auto logits_tensor = std::move(*logits_result); - - pos += 1; // start the loop from index 1 - start_pos += 1; - - while (pos < num_prompt_tokens) { - // Run the model - // NOLINTNEXTLINE(facebook-hte-ParameterUncheckedArrayBounds) - cur_token = prompt_tokens[pos]; - - auto step_result = text_decoder_runner_->step(tokens, start_pos); - if (!step_result.ok()) { - return step_result.error(); - } - logits_tensor = std::move(*step_result); - - pos++; - start_pos++; + ET_CHECK_OR_RETURN_ERROR( + use_kv_cache_ && decode_runner_ && + decode_runner_ != text_decoder_runner_, + OutOfResources, + "No prompt step fits the remaining context."); + // Choose decode before execution: an execution error is not retryable. + ET_CHECK_OK_OR_RETURN_ERROR(decode_runner_->load()); + } + auto* runner = + enable_parallel_prefill_ ? decode_runner_ : text_decoder_runner_; + uint64_t cur_token = 0; + auto tokens = + from_blob(&cur_token, {1, 1}, executorch::aten::ScalarType::Long); + for (size_t i = 0; i < prompt_tokens.size(); ++i) { + cur_token = prompt_tokens[i]; + auto logits = runner->step(tokens, start_pos); + ET_CHECK_OK_OR_RETURN_ERROR(logits.error()); + ++start_pos; + if (i + 1 == prompt_tokens.size()) { + cur_token = runner->logits_to_token(*logits, temperature); } - - cur_token = - text_decoder_runner_->logits_to_token(logits_tensor, temperature); } return cur_token; } diff --git a/extension/llm/runner/text_prefiller.h b/extension/llm/runner/text_prefiller.h index 9f032122349..d61ed466e07 100644 --- a/extension/llm/runner/text_prefiller.h +++ b/extension/llm/runner/text_prefiller.h @@ -23,7 +23,8 @@ class ET_EXPERIMENTAL TextPrefiller { TextDecoderRunner* text_decoder_runner, bool use_kv_cache, bool enable_parallel_prefill, - int64_t max_seq_len = 128); + int64_t max_seq_len = 128, + TextDecoderRunner* decode_runner = nullptr); virtual ~TextPrefiller() = default; /** @@ -94,6 +95,7 @@ class ET_EXPERIMENTAL TextPrefiller { * and passes the TextDecoderRunner instance to TextPrefiller. */ TextDecoderRunner* text_decoder_runner_; + TextDecoderRunner* decode_runner_; bool use_kv_cache_; bool enable_parallel_prefill_; int64_t max_seq_len_; diff --git a/shim_et/xplat/executorch/build/build_variables.bzl b/shim_et/xplat/executorch/build/build_variables.bzl index faeacb76b1d..20e2fc49040 100644 --- a/shim_et/xplat/executorch/build/build_variables.bzl +++ b/shim_et/xplat/executorch/build/build_variables.bzl @@ -364,6 +364,8 @@ EXTENSION_RUNNER_UTIL_SRCS = [ ] EXTENSION_LLM_RUNNER_SRCS = [ + "extension/llm/runner/io_manager/hf_static_io_manager.cpp", + "extension/llm/runner/io_manager/kv_cache_io_manager.cpp", "extension/llm/runner/llm_runner_helper.cpp", "extension/llm/runner/multimodal_prefiller.cpp", "extension/llm/runner/multimodal_runner.cpp", diff --git a/tools/cmake/executorch-config.cmake b/tools/cmake/executorch-config.cmake index 6ae9b2c12f3..aefb23d7f74 100644 --- a/tools/cmake/executorch-config.cmake +++ b/tools/cmake/executorch-config.cmake @@ -101,6 +101,7 @@ set(optional_lib_list metal_backend neuron_backend qnn_executorch_backend + qnn_llm_io_buffer_allocator portable_ops_lib custom_ops extension_asr_runner