From 1b88bb99571ef04dec7c71abd99fab770c059fd1 Mon Sep 17 00:00:00 2001 From: Siddartha Pothapragada Date: Tue, 3 Nov 2026 12:23:46 -0800 Subject: [PATCH 1/3] Add backend-contributed edge transform passes to Partitioner MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Run each partitioner’s edge transforms after caller-supplied passes so backends can use the standard lowering entry point. QNN contributes its existing transforms through the new hook. --- .../qualcomm/partition/qnn_partitioner.py | 37 ++++++- backends/qualcomm/tests/test_passes.py | 69 +++++++++++- backends/qualcomm/utils/utils.py | 32 ++---- exir/backend/partitioner.py | 26 +++++ exir/program/_program.py | 66 ++++++++++- exir/tests/targets.bzl | 14 +++ .../test_partitioner_transform_passes.py | 104 ++++++++++++++++++ 7 files changed, 320 insertions(+), 28 deletions(-) create mode 100644 exir/tests/test_partitioner_transform_passes.py diff --git a/backends/qualcomm/partition/qnn_partitioner.py b/backends/qualcomm/partition/qnn_partitioner.py index 23016de006b..b27b2eb94a3 100644 --- a/backends/qualcomm/partition/qnn_partitioner.py +++ b/backends/qualcomm/partition/qnn_partitioner.py @@ -5,10 +5,13 @@ # LICENSE file in the root directory of this source tree. import copy 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 +44,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 +203,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 +213,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 +235,26 @@ 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 + 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/tests/test_passes.py b/backends/qualcomm/tests/test_passes.py index 606fed01d72..235d3975006 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,70 @@ 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_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/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/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..a3889a62572 100644 --- a/exir/program/_program.py +++ b/exir/program/_program.py @@ -1131,12 +1131,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 +1158,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 +1296,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 +1414,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([])) From 33eee7617c51d02f61c94a42d5c10670c73c37c1 Mon Sep 17 00:00:00 2001 From: Siddartha Pothapragada Date: Tue, 3 Nov 2026 12:23:46 -0800 Subject: [PATCH 2/3] Avoid reused names when lifting constant tensors Reserve names for constant tensors already present in the exported program before lifting new constants, so methods with folded tables retain distinct constants. --- exir/program/_program.py | 12 +++++++++++- exir/tests/test_passes.py | 25 +++++++++++++++++++++++++ 2 files changed, 36 insertions(+), 1 deletion(-) diff --git a/exir/program/_program.py b/exir/program/_program.py index a3889a62572..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 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): From 3823fafc0816d68637e91cc2379700aa9b352cef Mon Sep 17 00:00:00 2001 From: Siddartha Pothapragada Date: Tue, 3 Nov 2026 12:51:52 -0800 Subject: [PATCH 3/3] LLM runner: select I/O contracts and accept custom buffer allocators Choose the registered IOManager and its logits, method names and step sizes, rejecting unknown contracts and invalid metadata. Load and reset the selected methods; route prompt tails to decode with a side-effect-free capacity check. Preserve module and I/O-buffer lifetimes and expose a backend-neutral allocator hook. --- extension/llm/runner/constants.h | 3 + extension/llm/runner/io_manager/io_manager.h | 156 +++++ .../runner/io_manager/kv_cache_io_manager.cpp | 555 ++++++++++++++++++ .../runner/io_manager/kv_cache_io_manager.h | 156 +++++ extension/llm/runner/io_manager/targets.bzl | 15 + extension/llm/runner/llm_runner_helper.cpp | 80 ++- extension/llm/runner/targets.bzl | 1 + extension/llm/runner/test/CMakeLists.txt | 10 +- extension/llm/runner/test/targets.bzl | 8 + extension/llm/runner/test/test_io_manager.cpp | 57 ++ .../llm/runner/test/test_text_llm_runner.cpp | 51 +- .../llm/runner/test/test_text_prefiller.cpp | 57 ++ extension/llm/runner/text_decoder_runner.cpp | 12 +- extension/llm/runner/text_decoder_runner.h | 7 + extension/llm/runner/text_llm_runner.cpp | 18 +- extension/llm/runner/text_llm_runner.h | 4 +- extension/llm/runner/text_prefiller.cpp | 86 ++- extension/llm/runner/text_prefiller.h | 4 +- .../executorch/build/build_variables.bzl | 1 + 19 files changed, 1207 insertions(+), 74 deletions(-) create mode 100644 extension/llm/runner/io_manager/kv_cache_io_manager.cpp create mode 100644 extension/llm/runner/io_manager/kv_cache_io_manager.h create mode 100644 extension/llm/runner/test/test_io_manager.cpp 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/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..0c230f2ac88 --- /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_builtin_io_managers() { + 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..79a350301bd --- /dev/null +++ b/extension/llm/runner/io_manager/kv_cache_io_manager.h @@ -0,0 +1,156 @@ +/* + * Copyright (c) Meta Platforms, Inc. and affiliates. + * All rights reserved. + * + * This source code is licensed under the BSD-style license found in the + * LICENSE file in the root directory of this source tree. + */ + +#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; +}; + +/// Registers the IOManagers this library implements (e.g. "kvcache_io") with +/// `io_manager_registry()`. Idempotent. +void register_builtin_io_managers(); + +} // 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..cd77ff1c64b 100644 --- a/extension/llm/runner/io_manager/targets.bzl +++ b/extension/llm/runner/io_manager/targets.bzl @@ -17,3 +17,18 @@ def define_common_targets(): ], 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..b92017c1152 100644 --- a/extension/llm/runner/llm_runner_helper.cpp +++ b/extension/llm/runner/llm_runner_helper.cpp @@ -10,6 +10,8 @@ #include #include + +#include #include #include #include @@ -26,6 +28,8 @@ #include #include +#include + namespace executorch::extension::llm { using ::executorch::extension::Module; @@ -126,6 +130,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 +159,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 +286,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 +333,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 +357,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 +379,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..28f71881083 100644 --- a/extension/llm/runner/targets.bzl +++ b/extension/llm/runner/targets.bzl @@ -161,6 +161,7 @@ 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: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/CMakeLists.txt b/extension/llm/runner/test/CMakeLists.txt index cd380429601..41175b3ee09 100644 --- a/extension/llm/runner/test/CMakeLists.txt +++ b/extension/llm/runner/test/CMakeLists.txt @@ -19,6 +19,7 @@ include(${EXECUTORCH_ROOT}/tools/cmake/Test.cmake) set(_test_srcs test_generation_config.cpp + test_io_manager.cpp test_model_metadata.cpp test_text_llm_runner.cpp test_text_prefiller.cpp @@ -49,6 +50,13 @@ 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 @@ -56,7 +64,7 @@ add_custom_command( # 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 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_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..97447fff8ed 100644 --- a/shim_et/xplat/executorch/build/build_variables.bzl +++ b/shim_et/xplat/executorch/build/build_variables.bzl @@ -364,6 +364,7 @@ EXTENSION_RUNNER_UTIL_SRCS = [ ] EXTENSION_LLM_RUNNER_SRCS = [ + "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",