diff --git a/backends/native/examples/hf_llm/runtime/BUCK b/backends/native/examples/hf_llm/runtime/BUCK index d4621cdb274..61772593fde 100644 --- a/backends/native/examples/hf_llm/runtime/BUCK +++ b/backends/native/examples/hf_llm/runtime/BUCK @@ -30,6 +30,7 @@ runtime.cxx_library( deps = [ ":llm_runner_utils", "//executorch/backends/native/examples/runner_utils:runner_utils", + "//executorch/backends/native/runtime:offgraph_cache", "//executorch/backends/native/runtime:runtime", "//executorch/backends/native/runtime/deserialize:package", "//executorch/backends/native/runtime/deserialize:samples", @@ -37,7 +38,6 @@ runtime.cxx_library( "//executorch/backends/native/runtime/graph:scalar_type", "//executorch/examples/models/llama/tokenizer:tiktoken", "//executorch/extension/llm/cache:kv_cache", - "//executorch/extension/llm/runner:constants", ], exported_deps = ["//executorch/backends/native/runtime/engine:engine"], external_deps = ["gflags"], diff --git a/backends/native/examples/hf_llm/runtime/LlmRunner.cpp b/backends/native/examples/hf_llm/runtime/LlmRunner.cpp index c7cfce656a8..4ee7a854c4d 100644 --- a/backends/native/examples/hf_llm/runtime/LlmRunner.cpp +++ b/backends/native/examples/hf_llm/runtime/LlmRunner.cpp @@ -14,7 +14,6 @@ #include #include #include -#include #include #include #include @@ -27,6 +26,7 @@ #include #include #include +#include #include #include #include @@ -35,7 +35,6 @@ #include #include #include -#include #include #include #include @@ -124,111 +123,6 @@ size_t kv_cache_length(const ptn::Method& method) { return length; } -constexpr std::string_view kUpdateAndAttend = - "torch.ops.kvcache.update_and_attend.default"; - -// The off-graph KV cache a method's update_and_attend calls need: one layer per -// cache id, shaped like its K, with the sliding window (or none) the package's -// metadata gives it. nullopt if the method has none. -std::optional offgraph_cache_geometry( - const ptn::Method& method, - const ptn::Package& package) { - std::map layers; - for (const ptn::Node& node : method.graph.nodes) { - if (node.target != kUpdateAndAttend) { - continue; - } - const auto arg = [&](std::string_view name, - size_t positional) -> const ptn::Argument& { - for (const ptn::NamedArgument& input : node.inputs) { - if (input.name == name) { - return input.arg; - } - } - if (positional >= node.inputs.size()) { - throw std::runtime_error( - "update_and_attend is missing argument " + std::string(name)); - } - return node.inputs[positional].arg; - }; - const ptn::Argument& layer_arg = arg("layer_id", 4); - const ptn::Argument& k_arg = arg("k", 1); - if (layer_arg.kind() != ptn::ArgKind::Int || - ptn::valid(layer_arg.as_int().id) || - k_arg.kind() != ptn::ArgKind::Tensor) { - throw std::runtime_error( - "update_and_attend needs a constant layer_id and a K tensor"); - } - const std::vector& k_sizes = - method.graph.value(k_arg.as_tensor().id).tensor_meta().sizes; - if (k_sizes.size() != 4) { - throw std::runtime_error("update_and_attend K must be [B, H, T, D]"); - } - const cache::LayerGeometry layer{ - .n_kv_heads = static_cast(k_sizes[1]), - .head_dim = static_cast(k_sizes[3])}; - const auto [it, inserted] = - layers.try_emplace(layer_arg.as_int().value, layer); - if (!inserted && - (it->second.n_kv_heads != layer.n_kv_heads || - it->second.head_dim != layer.head_dim)) { - throw std::runtime_error( - "update_and_attend calls disagree on cache layer " + - std::to_string(it->first) + "'s shape"); - } - } - if (layers.empty()) { - return std::nullopt; - } - cache::CacheGeometry geometry; - for (const auto& [id, layer] : layers) { - if (id != static_cast(geometry.layers.size())) { - throw std::runtime_error("update_and_attend cache ids must be 0..n-1"); - } - geometry.layers.push_back(layer); - } - - namespace llm = executorch::extension::llm; - const auto per_layer = [&](const char* name) { - std::optional> values = package.metadata_ints(name); - if (values && values->size() != geometry.layers.size()) { - throw std::runtime_error( - std::string("package metadata ") + name + " has " + - std::to_string(values->size()) + " entries for " + - std::to_string(geometry.layers.size()) + " cache layers"); - } - return values; - }; - const auto kv_heads = per_layer(llm::kKVHeads); - const auto head_dims = per_layer(llm::kHeadDims); - const auto windows = per_layer(llm::kWindows); - // Without them a sliding layer would silently attend its whole history. - if (!windows) { - throw std::runtime_error( - std::string("package metadata has no ") + llm::kWindows + - "; export the package with export_hf_llm.py --offgraph_cache"); - } - for (size_t i = 0; i < geometry.layers.size(); ++i) { - cache::LayerGeometry& layer = geometry.layers[i]; - if ((kv_heads && (*kv_heads)[i] != layer.n_kv_heads) || - (head_dims && (*head_dims)[i] != layer.head_dim)) { - throw std::runtime_error( - "package metadata disagrees with the graph on cache layer " + - std::to_string(i) + "'s shape"); - } - const int64_t window = (*windows)[i]; - if (window < 0 || window > std::numeric_limits::max()) { - throw std::runtime_error( - "package metadata gives cache layer " + std::to_string(i) + - " an invalid window"); - } - if (window > 0) { - layer.policy = {cache::LayerPolicy::Kind::Ring, static_cast(window)}; - } - } - return geometry; -} - // A runner-owned off-graph KV cache, published for the engine to bind. struct OffGraphCache { std::shared_ptr cache; @@ -605,16 +499,9 @@ int run_llm( const std::unique_ptr host = make_host(); // Declared before the context, which may hold the cache, so it outlives it. std::optional kv; - if (const auto geometry = offgraph_cache_geometry(method, *package)) { - // Serialized input sizes are upper bounds; the tokens are [1, T]. - const std::vector& token_bound = - method.graph.value(method.graph.input_ids.at(0)).tensor_meta().sizes; - if (token_bound.empty() || token_bound.back() <= 0 || - token_bound.back() > std::numeric_limits::max()) { - throw std::runtime_error("token input must have shape [1, sequence]"); - } + if (const auto geometry = ptn::offgraph_cache_geometry(method, *package)) { kv = install_offgraph_cache( - cache_backend_id, *geometry, static_cast(token_bound.back())); + cache_backend_id, *geometry, ptn::max_step_tokens(method)); } else if (FLAGS_interactive) { throw std::invalid_argument( "--interactive needs a package exported with --offgraph_cache"); diff --git a/backends/native/runtime/OffGraphCache.h b/backends/native/runtime/OffGraphCache.h new file mode 100644 index 00000000000..a08a30154af --- /dev/null +++ b/backends/native/runtime/OffGraphCache.h @@ -0,0 +1,145 @@ +// 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 + +// What a caller needs to build the off-graph KV cache a method's +// kvcache::update_and_attend calls use, read from the method and its package. + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +namespace ptn { + +// The cache geometry a method's update_and_attend calls need: one layer per +// cache id, shaped like its K, with the sliding window (or none) the package's +// metadata gives it. nullopt if the method makes no such calls. Throws +// std::runtime_error when the calls or the metadata disagree, or the metadata +// has no windows. +inline std::optional +offgraph_cache_geometry(const Method& method, const Package& package) { + namespace cache = executorch::extension::llm::cache; + namespace llm = executorch::extension::llm; + constexpr std::string_view kUpdateAndAttend = + "torch.ops.kvcache.update_and_attend.default"; + std::map layers; + for (const Node& node : method.graph.nodes) { + if (node.target != kUpdateAndAttend) { + continue; + } + const auto arg = [&](std::string_view name, + size_t positional) -> const Argument& { + for (const NamedArgument& input : node.inputs) { + if (input.name == name) { + return input.arg; + } + } + if (positional >= node.inputs.size()) { + throw std::runtime_error( + "update_and_attend is missing argument " + std::string(name)); + } + return node.inputs[positional].arg; + }; + const Argument& layer_arg = arg("layer_id", 4); + const Argument& k_arg = arg("k", 1); + if (layer_arg.kind() != ArgKind::Int || valid(layer_arg.as_int().id) || + k_arg.kind() != ArgKind::Tensor) { + throw std::runtime_error( + "update_and_attend needs a constant layer_id and a K tensor"); + } + const std::vector& k_sizes = + method.graph.value(k_arg.as_tensor().id).tensor_meta().sizes; + if (k_sizes.size() != 4) { + throw std::runtime_error("update_and_attend K must be [B, H, T, D]"); + } + const cache::LayerGeometry layer{ + .n_kv_heads = static_cast(k_sizes[1]), + .head_dim = static_cast(k_sizes[3])}; + const auto [it, inserted] = + layers.try_emplace(layer_arg.as_int().value, layer); + if (!inserted && + (it->second.n_kv_heads != layer.n_kv_heads || + it->second.head_dim != layer.head_dim)) { + throw std::runtime_error( + "update_and_attend calls disagree on cache layer " + + std::to_string(it->first) + "'s shape"); + } + } + if (layers.empty()) { + return std::nullopt; + } + cache::CacheGeometry geometry; + for (const auto& [id, layer] : layers) { + if (id != static_cast(geometry.layers.size())) { + throw std::runtime_error("update_and_attend cache ids must be 0..n-1"); + } + geometry.layers.push_back(layer); + } + + const auto per_layer = [&](const char* name) { + std::optional> values = package.metadata_ints(name); + if (values && values->size() != geometry.layers.size()) { + throw std::runtime_error( + std::string("package metadata ") + name + " has " + + std::to_string(values->size()) + " entries for " + + std::to_string(geometry.layers.size()) + " cache layers"); + } + return values; + }; + const auto kv_heads = per_layer(llm::kKVHeads); + const auto head_dims = per_layer(llm::kHeadDims); + const auto windows = per_layer(llm::kWindows); + // Without them a sliding layer would silently attend its whole history. + if (!windows) { + throw std::runtime_error( + std::string("package metadata has no ") + llm::kWindows + + "; export the package with export_hf_llm.py --offgraph_cache"); + } + for (size_t i = 0; i < geometry.layers.size(); ++i) { + cache::LayerGeometry& layer = geometry.layers[i]; + if ((kv_heads && (*kv_heads)[i] != layer.n_kv_heads) || + (head_dims && (*head_dims)[i] != layer.head_dim)) { + throw std::runtime_error( + "package metadata disagrees with the graph on cache layer " + + std::to_string(i) + "'s shape"); + } + const int64_t window = (*windows)[i]; + if (window < 0 || window > std::numeric_limits::max()) { + throw std::runtime_error( + "package metadata gives cache layer " + std::to_string(i) + + " an invalid window"); + } + if (window > 0) { + layer.policy = {cache::LayerPolicy::Kind::Ring, static_cast(window)}; + } + } + return geometry; +} + +// The most tokens one call of `method` takes, from its [1, T] token input's +// serialized upper bound: the max_write that sizes a ring layer. +inline int max_step_tokens(const Method& method) { + const std::vector& bound = + method.graph.value(method.graph.input_ids.at(0)).tensor_meta().sizes; + if (bound.empty() || bound.back() <= 0 || + bound.back() > std::numeric_limits::max()) { + throw std::runtime_error("token input must have shape [1, sequence]"); + } + return static_cast(bound.back()); +} + +} // namespace ptn diff --git a/backends/native/runtime/targets.bzl b/backends/native/runtime/targets.bzl index e95623ac14e..33d7dcdf526 100644 --- a/backends/native/runtime/targets.bzl +++ b/backends/native/runtime/targets.bzl @@ -150,6 +150,22 @@ def define_common_targets(): visibility = ["//executorch/backends/native/..."], ) + # Header-only: the off-graph KV cache a method's update_and_attend calls + # need, for the LLM runners and executors that build it. + runtime.cxx_library( + name = "offgraph_cache", + exported_headers = [ + "OffGraphCache.h", + ], + exported_deps = [ + ":method", + "//executorch/backends/native/runtime/deserialize:package", + "//executorch/extension/llm/cache:kv_cache", + "//executorch/extension/llm/runner:constants", + ], + visibility = ["PUBLIC"], + ) + runtime.cxx_library( name = "tensor_info", srcs = ["TensorInfo.cpp"],