From 80c68e909b87a0d48d097bc9255f21b7a4f297b2 Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Wed, 4 Nov 2026 18:08:55 -0800 Subject: [PATCH] Serve kvcache::update_and_attend in the native ggml engine The ggml engine now runs graphs exported with an off-graph KV cache. - GgmlSequenceCache is the ggml SequenceCache: the neutral sequence controller plus K/V storage per layer in the engine's backend buffer. Storage starts at the config's initial_capacity rows and doubles toward capacity, so short sequences do not pay for the full context. - register_ggml_cache_builders() registers it with the CacheFactory under backend id NativeGgml, kind single. - A context created with the llm_cache_registry_key option resolves the cache from the CacheRegistry. - update_and_attend lowers to ggml_set_rows into the cache and flash_attn_ext over the filled prefix, with per-step rows and causal mask uploads. Before each call the executor checks positions, grows the cache, rebinds storage (dropping cached graphs on a regrow), and plans every layer; it commits after the call succeeds. Flat layers, batch 1, and fp32/fp16 storage only. --- CMakeLists.txt | 6 + .../native/runtime/engine/ggml/CMakeLists.txt | 3 +- .../native/runtime/engine/ggml/GgmlCache.cpp | 181 ++++++++ .../native/runtime/engine/ggml/GgmlCache.h | 78 ++++ .../native/runtime/engine/ggml/GgmlEngine.cpp | 51 ++- .../native/runtime/engine/ggml/GgmlEngine.h | 23 +- .../runtime/engine/ggml/GgmlExecutor.cpp | 270 ++++++++++-- .../native/runtime/engine/ggml/GgmlExecutor.h | 7 +- .../runtime/engine/ggml/GgmlGraphBuilder.cpp | 6 +- .../runtime/engine/ggml/GgmlGraphBuilder.h | 25 +- .../runtime/engine/ggml/GgmlHostInts.cpp | 40 ++ .../engine/ggml/GgmlLowerAttention.cpp | 103 +++++ backends/native/runtime/engine/targets.bzl | 4 + backends/native/test/CMakeLists.txt | 2 +- backends/native/test/runtime/ggml/BUCK | 1 + .../test/runtime/ggml/test_ggml_engine.cpp | 388 +++++++++++++++++- 16 files changed, 1145 insertions(+), 43 deletions(-) create mode 100644 backends/native/runtime/engine/ggml/GgmlCache.cpp create mode 100644 backends/native/runtime/engine/ggml/GgmlCache.h diff --git a/CMakeLists.txt b/CMakeLists.txt index 08de715cc43..c4827241344 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -1045,6 +1045,12 @@ if(EXECUTORCH_BUILD_EXTENSION_RUNNER_UTIL) list(APPEND _executorch_extensions extension_runner_util) endif() +# The native ggml engine serves off-graph KV caches from the cache registry. +if(EXECUTORCH_BUILD_NATIVE_GGML AND NOT TARGET extension_llm_cache) + add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/llm/cache) + list(APPEND _executorch_extensions extension_llm_cache) +endif() + if(EXECUTORCH_BUILD_EXTENSION_TENSOR) add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/tensor) list(APPEND _executorch_extensions extension_tensor) diff --git a/backends/native/runtime/engine/ggml/CMakeLists.txt b/backends/native/runtime/engine/ggml/CMakeLists.txt index 8f8b7c07429..0df3d8fb80f 100644 --- a/backends/native/runtime/engine/ggml/CMakeLists.txt +++ b/backends/native/runtime/engine/ggml/CMakeLists.txt @@ -32,6 +32,7 @@ target_link_libraries(executorch_ggml INTERFACE ggml) add_library( native_ggml_engine + GgmlCache.cpp GgmlEngine.cpp GgmlExecutor.cpp GgmlGraphBuilder.cpp @@ -47,7 +48,7 @@ add_library(executorch::native_ggml_engine ALIAS native_ggml_engine) target_link_libraries( native_ggml_engine PUBLIC native_engine - PRIVATE executorch_ggml + PRIVATE executorch_ggml extension_llm_cache ) target_compile_options( native_ggml_engine PRIVATE $<$:/EHsc> diff --git a/backends/native/runtime/engine/ggml/GgmlCache.cpp b/backends/native/runtime/engine/ggml/GgmlCache.cpp new file mode 100644 index 00000000000..d3aeab3556b --- /dev/null +++ b/backends/native/runtime/engine/ggml/GgmlCache.cpp @@ -0,0 +1,181 @@ +// 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 + +#include +#include +#include +#include + +namespace ptn { + +namespace cache = executorch::extension::llm::cache; + +namespace { + +ggml_type kv_type(int kv_dtype) { + switch (static_cast(kv_dtype)) { + case ScalarType::Float: + return GGML_TYPE_F32; + case ScalarType::Half: + return GGML_TYPE_F16; + default: + throw std::invalid_argument( + "ggml cache: kv_dtype must be Float or Half, got " + + std::to_string(kv_dtype)); + } +} + +const cache::CacheGeometry& checked( + const cache::CacheGeometry& geometry, + const cache::CacheConfig& cfg) { + if (!cache::valid(geometry, cfg)) { + throw std::invalid_argument("ggml cache: invalid geometry or config"); + } + if (geometry.dflash_block_size) { + throw std::invalid_argument("ggml cache: DFlash geometry is unsupported"); + } + for (const cache::LayerGeometry& layer : geometry.layers) { + if (layer.policy.kind != cache::LayerPolicy::Kind::Flat) { + throw std::invalid_argument( + "ggml cache: only flat (full-history) layers are supported"); + } + } + return geometry; +} + +} // namespace + +GgmlSequenceCache::GgmlSequenceCache( + const cache::CacheGeometry& geometry, + const cache::CacheConfig& cfg) + : SequenceCache(checked(geometry, cfg), cfg), + geometry_(geometry), + type_(kv_type(cfg.kv_dtype)), + initial_rows_(std::clamp(cfg.initial_capacity, 1, cfg.capacity)) {} + +void GgmlSequenceCache::reserve(ggml_backend_t backend, int rows) { + ggml_backend_buffer_type_t buffer_type = + ggml_backend_get_default_buffer_type(backend); + if (storage_ && buffer_type != buffer_type_) { + throw std::runtime_error( + std::string("ggml cache: storage lives in ") + + ggml_backend_buft_name(buffer_type_) + ", not " + + ggml_backend_buft_name(buffer_type)); + } + if (rows > capacity()) { + throw std::out_of_range( + "ggml cache: " + std::to_string(rows) + + " rows exceed the capacity of " + std::to_string(capacity())); + } + if (storage_ && rows <= storage_->rows) { + return; + } + int64_t target = storage_ ? storage_->rows : initial_rows_; + while (target < rows) { + target *= 2; + } + auto next = std::make_shared(); + next->rows = static_cast(std::min(target, capacity())); + + ggml_init_params params{ + .mem_size = ggml_tensor_overhead() * (2 * geometry_.layers.size() + 1), + .mem_buffer = nullptr, + .no_alloc = true, + }; + next->context.reset(ggml_init(params)); + if (!next->context) { + throw std::runtime_error("ggml cache: failed to create tensor context"); + } + for (size_t i = 0; i < geometry_.layers.size(); ++i) { + const cache::LayerGeometry& layer = geometry_.layers[i]; + Layer out; + for (ggml_tensor** t : {&out.k, &out.v}) { + *t = ggml_new_tensor_3d( + next->context.get(), + type_, + layer.head_dim, + next->rows, + layer.n_kv_heads); + ggml_format_name(*t, "kv_cache.%s.%zu", t == &out.k ? "k" : "v", i); + } + next->layers.push_back(out); + } + next->buffer.reset( + ggml_backend_alloc_ctx_tensors_from_buft(next->context.get(), buffer_type)); + if (!next->buffer) { + throw std::runtime_error( + "ggml cache: failed to allocate K/V storage of " + + std::to_string(next->rows) + " rows"); + } + ggml_backend_buffer_clear(next->buffer.get(), 0); + + // Each head's rows are contiguous, but the head stride grows with the rows. + std::vector staging; + for (size_t i = 0; storage_ && i < storage_->layers.size(); ++i) { + const Layer& old = storage_->layers[i]; + for (const auto& [from, to] : + {std::pair{old.k, next->layers[i].k}, + std::pair{old.v, next->layers[i].v}}) { + staging.resize(ggml_nbytes(from)); + ggml_backend_tensor_get(from, staging.data(), 0, staging.size()); + const size_t head_bytes = static_cast(from->ne[1]) * from->nb[1]; + for (int64_t h = 0; h < from->ne[2]; ++h) { + ggml_backend_tensor_set( + to, + staging.data() + h * from->nb[2], + h * to->nb[2], + head_bytes); + } + } + } + + buffer_type_ = buffer_type; + storage_ = std::move(next); +} + +void* GgmlSequenceCache::face(cache::FaceId id) { + return cache::expose< + cache::SequenceControl, + cache::SequencePlanner, + GgmlSequenceCache>(this, id); +} + +namespace { + +// CacheFactory::build reports a null cache as an error, so a cache the engine +// cannot serve comes back as one rather than as an exception. +template +cache::CacheBuilder builder() { + return [](const cache::CacheGeometry& geometry, + const cache::CacheConfig& cfg) -> std::shared_ptr { + try { + return std::make_shared(geometry, cfg); + } catch (const std::exception& error) { + ET_LOG(Error, "%s", error.what()); + return nullptr; + } + }; +} + +} // namespace + +void register_ggml_cache_builders() { + // Registering twice is harmless; the first builder stays. + (void)cache::CacheFactory::global().register_builder( + kGgmlCacheBackendId, cache::kind::kSingle, builder()); +} + +} // namespace ptn diff --git a/backends/native/runtime/engine/ggml/GgmlCache.h b/backends/native/runtime/engine/ggml/GgmlCache.h new file mode 100644 index 00000000000..8f5e46c0e21 --- /dev/null +++ b/backends/native/runtime/engine/ggml/GgmlCache.h @@ -0,0 +1,78 @@ +// 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 +#include + +#include + +namespace ptn { + +// The off-graph KV cache of one sequence for the ggml engine: the neutral +// SequenceCache controller plus K/V storage per layer in the buffer type of +// the first backend that binds it. Each layer's K and V are ggml +// [head_dim, rows, n_kv_heads] (torch [1, n_kv_heads, rows, head_dim]); row r +// holds logical position r. Storage starts at the config's initial_capacity +// rows and doubles toward capacity as steps need more. Like SequenceCache, it +// is not thread-safe: executables bound to one cache must be serialized. +class GgmlSequenceCache final + : public executorch::extension::llm::cache::SequenceCache { + public: + static constexpr const char* kFaceName = "ptn.ggml.SequenceCache"; + + struct Layer { + ggml_tensor* k = nullptr; + ggml_tensor* v = nullptr; + }; + + // One allocation of K/V storage. A regrow moves the rows into a new one; + // the old one lives on while any executable still holds it. + struct Storage { + ggml_context_ptr context; + ggml_backend_buffer_ptr buffer; + std::vector layers; + int rows = 0; + }; + + // Throws on geometry or config the ggml engine cannot serve. + GgmlSequenceCache( + const executorch::extension::llm::cache::CacheGeometry& geometry, + const executorch::extension::llm::cache::CacheConfig& cfg); + + GgmlSequenceCache(const GgmlSequenceCache&) = delete; + GgmlSequenceCache& operator=(const GgmlSequenceCache&) = delete; + + // Ensures storage of at least `rows` rows in `backend`'s default buffer + // type, moving the rows it holds into a larger allocation if needed. Throws + // if `rows` exceeds capacity or an earlier call used another buffer type. + void reserve(ggml_backend_t backend, int rows); + + // The current storage; null before the first reserve. Graphs built over a + // storage stay valid while their executable holds it. + const std::shared_ptr& storage() const { + return storage_; + } + + protected: + void* face(executorch::extension::llm::cache::FaceId id) override; + + private: + executorch::extension::llm::cache::CacheGeometry geometry_; + ggml_type type_; + int initial_rows_; + ggml_backend_buffer_type_t buffer_type_ = nullptr; + std::shared_ptr storage_; +}; + +} // namespace ptn diff --git a/backends/native/runtime/engine/ggml/GgmlEngine.cpp b/backends/native/runtime/engine/ggml/GgmlEngine.cpp index 597445905ee..70429eb8c25 100644 --- a/backends/native/runtime/engine/ggml/GgmlEngine.cpp +++ b/backends/native/runtime/engine/ggml/GgmlEngine.cpp @@ -16,7 +16,9 @@ #include #include #include +#include #include +#include namespace ptn { @@ -81,17 +83,52 @@ std::unique_ptr GgmlEngineHost::create_context( return create_ggml_context(std::move(program), std::move(package)); } +std::unique_ptr GgmlEngineHost::create_context_with_options( + std::shared_ptr program, + std::shared_ptr package, + const EngineOptions& options) { + return create_ggml_context(std::move(program), std::move(package), options); +} + std::unique_ptr GgmlEngineHost::create_ggml_context( std::shared_ptr program, - std::shared_ptr package) { + std::shared_ptr package, + const EngineOptions& options) { + namespace cache = executorch::extension::llm::cache; + std::shared_ptr kv_cache; + for (const auto& [name, value] : options) { + if (name != cache::kCacheKeyOption) { + throw std::invalid_argument( + impl_->name + ": unsupported context option '" + name + "'"); + } + std::shared_ptr installed = + cache::CacheRegistry::global().get(value); + if (!installed) { + throw std::invalid_argument( + impl_->name + ": no KV cache is installed under key '" + value + "'"); + } + GgmlSequenceCache* ggml_cache = installed->as(); + if (ggml_cache == nullptr) { + throw std::invalid_argument( + impl_->name + ": KV cache '" + value + + "' was not built for this engine (CacheFactory backend '" + + kGgmlCacheBackendId + "', kind 'single')"); + } + kv_cache = std::shared_ptr(installed, ggml_cache); + } return std::unique_ptr(new GgmlEngineContext( - std::move(program), std::move(package), impl_->options)); + std::move(program), + std::move(package), + impl_->options, + std::move(kv_cache))); } struct GgmlEngineContext::Impl { - explicit Impl(GgmlEngineOptions value) : options(std::move(value)) {} + Impl(GgmlEngineOptions value, std::shared_ptr cache) + : options(std::move(value)), kv_cache(std::move(cache)) {} GgmlEngineOptions options; + std::shared_ptr kv_cache; size_t materialized_constant_bytes = 0; size_t heap_copied_constant_bytes = 0; GgmlMutableKeyOwners mutable_key_owners; @@ -100,9 +137,10 @@ struct GgmlEngineContext::Impl { GgmlEngineContext::GgmlEngineContext( std::shared_ptr program, std::shared_ptr package, - GgmlEngineOptions options) + GgmlEngineOptions options, + std::shared_ptr cache) : EngineContext(std::move(program), std::move(package)), - impl_(std::make_unique(std::move(options))) {} + impl_(std::make_unique(std::move(options), std::move(cache))) {} GgmlEngineContext::~GgmlEngineContext() = default; @@ -114,7 +152,8 @@ std::unique_ptr GgmlEngineContext::compile_method( impl_->options, &impl_->materialized_constant_bytes, &impl_->heap_copied_constant_bytes, - impl_->mutable_key_owners); + impl_->mutable_key_owners, + impl_->kv_cache.get()); } size_t GgmlEngineContext::materialized_constant_bytes() const { diff --git a/backends/native/runtime/engine/ggml/GgmlEngine.h b/backends/native/runtime/engine/ggml/GgmlEngine.h index 8fad5e69e17..0269baa79f5 100644 --- a/backends/native/runtime/engine/ggml/GgmlEngine.h +++ b/backends/native/runtime/engine/ggml/GgmlEngine.h @@ -14,6 +14,15 @@ namespace ptn { +// CacheFactory backend id of the ggml engine's off-graph KV caches. A method +// that calls kvcache::update_and_attend compiles only in a context created +// with the kCacheKeyOption ("llm_cache_registry_key") option naming a cache +// built under this id with kind::kSingle. +inline constexpr char kGgmlCacheBackendId[] = "NativeGgml"; + +// Registers the ggml cache builders with CacheFactory::global(). Idempotent. +void register_ggml_cache_builders(); + enum class GgmlDevice { Cpu, // The GPU backend is whichever one the linked ggml variant registers (CUDA, @@ -43,6 +52,7 @@ struct GgmlEngineOptions { }; class GgmlEngineContext; +class GgmlSequenceCache; class GgmlEngineHost final : public EngineHost { public: @@ -57,9 +67,17 @@ class GgmlEngineHost final : public EngineHost { std::shared_ptr program, std::shared_ptr package) override; + // Accepts kCacheKeyOption, resolved through CacheRegistry::global(). Contexts + // bound to one cache share its sequence: the caller must serialize them. + std::unique_ptr create_context_with_options( + std::shared_ptr program, + std::shared_ptr package, + const EngineOptions& options) override; + std::unique_ptr create_ggml_context( std::shared_ptr program, - std::shared_ptr package); + std::shared_ptr package, + const EngineOptions& options = {}); private: struct Impl; @@ -86,7 +104,8 @@ class GgmlEngineContext final : public EngineContext { GgmlEngineContext( std::shared_ptr program, std::shared_ptr package, - GgmlEngineOptions options); + GgmlEngineOptions options, + std::shared_ptr cache); std::unique_ptr compile_method( const Method& method) override; diff --git a/backends/native/runtime/engine/ggml/GgmlExecutor.cpp b/backends/native/runtime/engine/ggml/GgmlExecutor.cpp index 66e5c2cf02a..8866f77e75f 100644 --- a/backends/native/runtime/engine/ggml/GgmlExecutor.cpp +++ b/backends/native/runtime/engine/ggml/GgmlExecutor.cpp @@ -12,6 +12,7 @@ #include #include #include +#include #include #include #include @@ -33,6 +34,7 @@ #include #include #include +#include #include #include #include @@ -48,6 +50,7 @@ namespace ptn { namespace { using namespace ggml_detail; +using executorch::extension::llm::cache::SeqStepPlan; // Multiplies each slice of the leading dimension by its scale. // data is a byte buffer holding count floats, as ggml tensor data is. @@ -161,17 +164,20 @@ class GgmlExecutable final : public EngineExecutable { const GgmlEngineOptions& options, size_t* materialized_constant_bytes, size_t* heap_copied_constant_bytes, - GgmlMutableKeyOwners& mutable_key_owners) + GgmlMutableKeyOwners& mutable_key_owners, + GgmlSequenceCache* kv_cache) : mutable_key_claim_(mutable_key_owners, method), method_(method), package_(package), - options_(options) { + options_(options), + kv_cache_(kv_cache) { index_signature(); initialize_backends(); fold_batch_norms(); initialize_persistent_tensors(materialized_constant_bytes); index_rope_frequencies(); index_mask_extents(); + index_offgraph_attention(); if (heap_copied_constant_bytes != nullptr) { *heap_copied_constant_bytes += heap_copied_constant_bytes_; } @@ -256,27 +262,55 @@ class GgmlExecutable final : public EngineExecutable { })) { throw std::logic_error("ggml: every input must be set before execute"); } + if (valid(offgraph_position_)) { + offgraph_step_ = checked_offgraph_step(); + } - if (options_.max_cached_graphs == 0) { - ggml_backend_sched_reset(scheduler_.get()); - GraphInstance instance = build_graph(evaluate_key_extents()); - allocate(scheduler_.get(), instance); - try { - run(scheduler_.get(), instance); - } catch (...) { + // A step that fails before its graph runs leaves the cache unchanged; one + // that fails while running may have written rows from its start. + bool ran = false; + std::optional step; + try { + if (valid(offgraph_position_)) { + bind_offgraph(offgraph_step_->start + offgraph_step_->tokens); + } + if (options_.max_cached_graphs == 0) { ggml_backend_sched_reset(scheduler_.get()); - throw; + GraphInstance instance = build_graph(evaluate_key_extents()); + step = prepare_offgraph(instance); + allocate(scheduler_.get(), instance); + ran = true; + try { + run(scheduler_.get(), instance); + } catch (...) { + ggml_backend_sched_reset(scheduler_.get()); + throw; + } + ggml_backend_sched_reset(scheduler_.get()); + } else { + CachedGraph& cached = cached_graph(); + step = prepare_offgraph(cached.instance); + ran = true; + try { + run(cached.scheduler.get(), cached.instance); + } catch (...) { + ggml_backend_sched_synchronize(cached.scheduler.get()); + cached_graphs_.pop_back(); + throw; + } } - ggml_backend_sched_reset(scheduler_.get()); - } else { - CachedGraph& cached = cached_graph(); - try { - run(cached.scheduler.get(), cached.instance); - } catch (...) { - ggml_backend_sched_synchronize(cached.scheduler.get()); - cached_graphs_.pop_back(); - throw; + } catch (...) { + if (ran && valid(offgraph_position_) && + offgraph_step_->start < kv_cache_->length()) { + kv_cache_->rewind(static_cast(offgraph_step_->start)); } + throw; + } + if (step) { + if (offgraph_step_->start < kv_cache_->length()) { + kv_cache_->rewind(static_cast(offgraph_step_->start)); + } + kv_cache_->commit(*step); } executed_ = true; } @@ -1182,18 +1216,25 @@ class GgmlExecutable final : public EngineExecutable { std::vector full_key_extents() const { std::vector out(mask_extents_.size()); std::ranges::transform(mask_extents_, out.begin(), &MaskExtent::keys); + if (valid(offgraph_position_)) { + out.push_back(offgraph_storage_->rows); + } return out; } - std::vector evaluate_key_extents() { - if (mask_extents_.empty()) { - return full_key_extents(); - } - HostIntEvaluator host( + HostIntEvaluator host_evaluator() { + return HostIntEvaluator( method_, inputs_, [this](ValueId id) { return host_constant(id); }, host_static_values_); + } + + std::vector evaluate_key_extents() { + if (mask_extents_.empty() && !valid(offgraph_position_)) { + return full_key_extents(); + } + HostIntEvaluator host = host_evaluator(); std::vector out; for (const MaskExtent& mask : mask_extents_) { const HostInts* key_index = host.tensor(mask.key_operand); @@ -1203,9 +1244,166 @@ class GgmlExecutable final : public EngineExecutable { ? attended_key_extent(mask, *key_index, *query) : mask.keys); } + if (valid(offgraph_position_)) { + // The cache rows up to the step's last position, bucketed like masks. + const int64_t needed = offgraph_step_->start + offgraph_step_->tokens; + out.push_back(std::min( + offgraph_storage_->rows, + (needed + kKeyExtentBucket - 1) / kKeyExtentBucket * + kKeyExtentBucket)); + } return out; } + // Binds the off-graph cache to a method that calls update_and_attend. + void index_offgraph_attention() { + for (const NodeId id : method_.graph.schedule) { + const Node& node = method_.graph.node(id); + if (!node.is_call() || + op_key(node.target) != "kvcache.update_and_attend.default") { + continue; + } + const NamedArgument* position = find_arg(node, "position", 3); + if (position == nullptr || position->arg.kind() != ArgKind::Tensor) { + throw std::runtime_error( + "ggml: update_and_attend is missing its position tensor"); + } + const ValueId position_id = position->arg.as_tensor().id; + if (valid(offgraph_position_) && position_id != offgraph_position_) { + throw std::runtime_error( + "ggml: every update_and_attend in a method must read the same " + "position tensor"); + } + offgraph_position_ = position_id; + } + if (!valid(offgraph_position_)) { + return; + } + if (kv_cache_ == nullptr) { + throw std::runtime_error( + "ggml: method '" + method_.name + + "' calls kvcache.update_and_attend but its context has no KV " + "cache; create the context with the llm_cache_registry_key option"); + } + bind_offgraph(1); + } + + // Grows the cache to `rows` and, if its storage moved, rebinds it and drops + // the graphs built over the old storage, which this executable keeps alive + // until then. + void bind_offgraph(int64_t rows) { + kv_cache_->reserve(primary_backend_, static_cast(rows)); + if (kv_cache_->storage() == offgraph_storage_) { + return; + } + for (CachedGraph& cached : cached_graphs_) { + ggml_backend_sched_synchronize(cached.scheduler.get()); + } + cached_graphs_.clear(); + offgraph_storage_ = kv_cache_->storage(); + offgraph_cache_.layers.clear(); + for (const GgmlSequenceCache::Layer& layer : offgraph_storage_->layers) { + offgraph_cache_.layers.emplace_back(layer.k, layer.v); + } + } + + struct OffGraphStep { + int64_t start; + int64_t tokens; + }; + + // Reads the step's positions on the host and checks they are consecutive, + // fit the cache, and start no later than the sequence's length, at a + // position the sequence can rewind to. A step that starts earlier rewinds + // the sequence to its start once it has run. + OffGraphStep checked_offgraph_step() { + HostIntEvaluator host = host_evaluator(); + const HostInts* values = host.tensor(offgraph_position_); + if (values == nullptr || values->data.empty()) { + throw std::runtime_error( + "ggml: update_and_attend positions must be computable on the host " + "from graph inputs"); + } + const std::vector& positions = values->data; + const int64_t start = positions.front(); + for (size_t i = 0; i < positions.size(); ++i) { + if (positions[i] != start + static_cast(i)) { + throw std::invalid_argument( + "ggml: update_and_attend positions must be consecutive"); + } + } + const auto tokens = static_cast(positions.size()); + if (start < 0 || start + tokens > kv_cache_->capacity()) { + throw std::out_of_range( + "ggml: KV cache step [" + std::to_string(start) + ", " + + std::to_string(start + tokens) + ") exceeds the cache capacity of " + + std::to_string(kv_cache_->capacity())); + } + if (start > kv_cache_->length()) { + throw std::invalid_argument( + "ggml: KV cache step starts at position " + std::to_string(start) + + ", past the sequence's length of " + + std::to_string(kv_cache_->length())); + } + if (start < kv_cache_->length() && + !kv_cache_->can_rewind(static_cast(start))) { + throw std::runtime_error( + "ggml: the KV cache cannot rewind to position " + + std::to_string(start)); + } + return {start, tokens}; + } + + // Plans the step on the cache and fills the graph's cache rows and mask. + // Returns the plan to commit once the step has run. + std::optional prepare_offgraph(GraphInstance& instance) { + if (!instance.offgraph) { + return std::nullopt; + } + const OffGraphUploads& slots = *instance.offgraph; + if (offgraph_step_->tokens != slots.tokens) { + throw std::runtime_error( + "ggml: update_and_attend needs one position per token"); + } + const int64_t start = offgraph_step_->start; + std::optional plan; + for (size_t layer = 0; layer < offgraph_cache_.layers.size(); ++layer) { + plan = kv_cache_->plan( + static_cast(layer), + static_cast(start), + static_cast(slots.tokens)); + if (!plan || plan->n_write != 1 || plan->n_read != 1 || + plan->read_base_pos != 0 || plan->read[0].start != 0) { + throw std::runtime_error( + "ggml: the KV cache planned a non-flat step"); + } + } + const int64_t read_end = plan->read[0].len; + if (read_end > slots.keys) { + throw std::logic_error("ggml: KV cache read exceeds the graph's keys"); + } + std::vector& rows = instance.uploads[slots.rows].bytes; + rows.resize(static_cast(slots.tokens) * sizeof(int32_t)); + auto* row = reinterpret_cast(rows.data()); + for (int64_t i = 0; i < slots.tokens; ++i) { + row[i] = plan->write[0].start + static_cast(i); + } + std::vector& mask = instance.uploads[slots.mask].bytes; + mask.resize( + static_cast(slots.tokens * slots.keys) * sizeof(ggml_fp16_t)); + auto* bias = reinterpret_cast(mask.data()); + const ggml_fp16_t attend = ggml_fp32_to_fp16(0.0f); + const ggml_fp16_t masked = + ggml_fp32_to_fp16(-std::numeric_limits::infinity()); + for (int64_t i = 0; i < slots.tokens; ++i) { + for (int64_t j = 0; j < slots.keys; ++j) { + bias[i * slots.keys + j] = + j <= start + i && j < read_end ? attend : masked; + } + } + return plan; + } + // Small integer constants, e.g. a key offset, read once from the package. std::optional host_constant(ValueId id) { if (const auto found = host_constants_.find(id); @@ -1275,6 +1473,16 @@ class GgmlExecutable final : public EngineExecutable { const Method& method_; const Package& package_; GgmlEngineOptions options_; + GgmlSequenceCache* kv_cache_; + // The position tensor of the method's update_and_attend calls, if any. + ValueId offgraph_position_ = kInvalid; + // Cache storage; build_graph sets the keys per graph. + OffGraphCacheBinding offgraph_cache_; + // The storage offgraph_cache_ and cached graphs point into. Declared before + // cached_graphs_ so it outlives them. + std::shared_ptr offgraph_storage_; + // The step execute() is running, checked once. + std::optional offgraph_step_; std::vector inputs_; std::vector outputs_; std::vector backends_; @@ -1308,6 +1516,10 @@ GraphInstance GgmlExecutable::build_graph( } } } + OffGraphCacheBinding offgraph_cache = offgraph_cache_; + if (valid(offgraph_position_)) { + offgraph_cache.keys = key_extents.back(); + } GraphInstance instance = GraphBuilder( method_, inputs_, @@ -1316,7 +1528,9 @@ GraphInstance GgmlExecutable::build_graph( rope_freq_factors_, batch_norms_, std::move(truncated), - options_.graph_size) + options_.graph_size, + valid(offgraph_position_) ? &offgraph_cache + : nullptr) .build(); if (options_.print_graph_ops) { std::map counts; @@ -1360,14 +1574,16 @@ std::unique_ptr make_ggml_executable( const GgmlEngineOptions& options, size_t* materialized_constant_bytes, size_t* heap_copied_constant_bytes, - GgmlMutableKeyOwners& mutable_key_owners) { + GgmlMutableKeyOwners& mutable_key_owners, + GgmlSequenceCache* kv_cache) { return std::make_unique( method, package, options, materialized_constant_bytes, heap_copied_constant_bytes, - mutable_key_owners); + mutable_key_owners, + kv_cache); } } // namespace ptn diff --git a/backends/native/runtime/engine/ggml/GgmlExecutor.h b/backends/native/runtime/engine/ggml/GgmlExecutor.h index 34323f8584d..6b4505093fc 100644 --- a/backends/native/runtime/engine/ggml/GgmlExecutor.h +++ b/backends/native/runtime/engine/ggml/GgmlExecutor.h @@ -29,13 +29,16 @@ ggml_backend_dev_t select_ggml_gpu_device(const GgmlEngineOptions& options); using GgmlMutableKeyOwners = std::unordered_map; // The executable holds the method's mutable buffer keys in -// `mutable_key_owners` until it is destroyed. +// `mutable_key_owners` until it is destroyed. `kv_cache`, which must outlive +// the executable, serves the method's kvcache::update_and_attend calls; a +// method with such calls and a null cache does not compile. std::unique_ptr make_ggml_executable( const Method& method, const Package& package, const GgmlEngineOptions& options, size_t* materialized_constant_bytes, size_t* heap_copied_constant_bytes, - GgmlMutableKeyOwners& mutable_key_owners); + GgmlMutableKeyOwners& mutable_key_owners, + GgmlSequenceCache* kv_cache = nullptr); } // namespace ptn diff --git a/backends/native/runtime/engine/ggml/GgmlGraphBuilder.cpp b/backends/native/runtime/engine/ggml/GgmlGraphBuilder.cpp index ee6718ebacc..c5522173893 100644 --- a/backends/native/runtime/engine/ggml/GgmlGraphBuilder.cpp +++ b/backends/native/runtime/engine/ggml/GgmlGraphBuilder.cpp @@ -204,7 +204,8 @@ GraphBuilder::GraphBuilder( const std::unordered_map>& rope_factor_bytes, const FoldedBatchNorms& batch_norms, std::unordered_map key_extents, - size_t graph_size) + size_t graph_size, + const OffGraphCacheBinding* offgraph_cache) : method_(method), inputs_(inputs), outputs_(outputs), @@ -213,6 +214,7 @@ GraphBuilder::GraphBuilder( batch_norms_(batch_norms), key_extents_(std::move(key_extents)), graph_size_(graph_size), + offgraph_cache_(offgraph_cache), values_(method.graph.values.size()) { for (const auto& [id, tensor] : persistent) { persistent_storage_.insert(tensor); @@ -743,6 +745,8 @@ void GraphBuilder::lower(const Node& node) { lower_index_put(node); } else if (key == "aten.scaled_dot_product_attention.default") { lower_sdpa(node); + } else if (key == "kvcache.update_and_attend.default") { + lower_update_and_attend(node); } else if (key == "aten.sym_size.int") { lower_sym_size(node); } else if ( diff --git a/backends/native/runtime/engine/ggml/GgmlGraphBuilder.h b/backends/native/runtime/engine/ggml/GgmlGraphBuilder.h index 2f69e368fe0..114bea57e66 100644 --- a/backends/native/runtime/engine/ggml/GgmlGraphBuilder.h +++ b/backends/native/runtime/engine/ggml/GgmlGraphBuilder.h @@ -88,6 +88,25 @@ struct RowIndexCheck { int64_t rows = 0; }; +// The off-graph KV cache as one graph sees it. +struct OffGraphCacheBinding { + // Per cache layer: K and V storage, ggml [head_dim, capacity, n_kv_heads]. + std::vector> layers; + // Cache rows the graph's attention reads, a prefix of capacity. + int64_t keys = 0; +}; + +// Uploads of a graph's off-graph attention, which the executor fills from the +// cache's plan before every run. Every layer shares them. +struct OffGraphUploads { + // I32 [tokens]: the cache row each token writes. + size_t rows = 0; + // F16 [tokens, keys]: additive mask of the rows each token attends. + size_t mask = 0; + int64_t tokens = 0; + int64_t keys = 0; +}; + struct GraphInstance { ggml_context_ptr context; ggml_cgraph* graph = nullptr; @@ -95,6 +114,7 @@ struct GraphInstance { std::vector outputs; // Input-driven row indices, rechecked on every replay. std::vector row_index_checks; + std::optional offgraph; }; // Inference batch norm is the per-channel map x * scale + shift of its @@ -136,7 +156,8 @@ class GraphBuilder { rope_factor_bytes, const FoldedBatchNorms& batch_norms, std::unordered_map key_extents, - size_t graph_size); + size_t graph_size, + const OffGraphCacheBinding* offgraph_cache = nullptr); GraphInstance build(); @@ -267,6 +288,7 @@ class GraphBuilder { void lower_index_put(const Node& node); ggml_tensor* attention_mask(ValueId id, int64_t tokens, int64_t keys); void lower_sdpa(const Node& node); + void lower_update_and_attend(const Node& node); void lower_sym_size(const Node& node); int64_t symbolic_operand(const Argument& arg) const; void lower_symbolic(const Node& node, const std::string& key); @@ -280,6 +302,7 @@ class GraphBuilder { // Mask values (and their producers) computed over a key prefix only. std::unordered_map key_extents_; size_t graph_size_; + const OffGraphCacheBinding* offgraph_cache_; std::vector values_; std::unordered_map rope_factor_tensors_; std::map, ggml_tensor*> attention_masks_; diff --git a/backends/native/runtime/engine/ggml/GgmlHostInts.cpp b/backends/native/runtime/engine/ggml/GgmlHostInts.cpp index d41b5f99383..cefebd10122 100644 --- a/backends/native/runtime/engine/ggml/GgmlHostInts.cpp +++ b/backends/native/runtime/engine/ggml/GgmlHostInts.cpp @@ -248,6 +248,46 @@ std::optional HostIntEvaluator::compute(ValueId id) { out.sizes.insert(out.sizes.begin() + *dim, 1); return out; } + if (op == "aten.select.int" || op == "aten.select_copy.int") { + const HostInts* self = arg_tensor(node, "self", 0); + std::optional dim = arg_int(node, "dim", 1); + std::optional index = arg_int(node, "index", 2); + if (self == nullptr || !dim || !index || self->sizes.empty()) { + return std::nullopt; + } + const auto rank = static_cast(self->sizes.size()); + if (*dim < 0) { + *dim += rank; + } + if (*dim < 0 || *dim >= rank) { + return std::nullopt; + } + const int64_t extent = self->sizes[static_cast(*dim)]; + if (*index < 0) { + *index += extent; + } + if (*index < 0 || *index >= extent) { + return std::nullopt; + } + size_t inner = 1; + for (int64_t d = *dim + 1; d < rank; ++d) { + inner *= static_cast(self->sizes[static_cast(d)]); + } + const size_t outer = self->data.size() / (inner * static_cast(extent)); + HostInts out; + out.sizes = self->sizes; + out.sizes.erase(out.sizes.begin() + *dim); + for (size_t o = 0; o < outer; ++o) { + const size_t offset = + (o * static_cast(extent) + static_cast(*index)) * + inner; + const auto first = + self->data.begin() + static_cast(offset); + out.data.insert( + out.data.end(), first, first + static_cast(inner)); + } + return out; + } if (op == "aten.view.default" || op == "aten.view_copy.default" || op == "aten.reshape.default" || op == "aten._unsafe_view.default") { const HostInts* self = arg_tensor(node, "self", 0); diff --git a/backends/native/runtime/engine/ggml/GgmlLowerAttention.cpp b/backends/native/runtime/engine/ggml/GgmlLowerAttention.cpp index aa8a2a5bf40..3b00c31e966 100644 --- a/backends/native/runtime/engine/ggml/GgmlLowerAttention.cpp +++ b/backends/native/runtime/engine/ggml/GgmlLowerAttention.cpp @@ -397,4 +397,107 @@ void GraphBuilder::lower_sdpa(const Node& node) { {query.sizes[0], query.sizes[1], tokens, val.sizes[3]}}); } +// Writes this step's K/V into the cache rows the executor uploads, then +// attends the cache prefix under the uploaded causal mask. Q/K/V are [B,H,T,D]. +void GraphBuilder::lower_update_and_attend(const Node& node) { + if (offgraph_cache_ == nullptr) { + throw std::runtime_error( + "ggml: kvcache.update_and_attend needs an off-graph KV cache; create " + "the context with the llm_cache_registry_key option"); + } + const TensorRef query = rows_arg(node, "q", 0); + const TensorRef key = rows_arg(node, "k", 1); + const TensorRef val = rows_arg(node, "v", 2); + const int64_t layer = int_arg(node, "layer_id", 4, -1); + const auto num_layers = + static_cast(offgraph_cache_->layers.size()); + if (layer < 0 || layer >= num_layers) { + throw std::runtime_error( + "ggml: update_and_attend layer_id " + std::to_string(layer) + + " is outside the cache's " + std::to_string(num_layers) + " layers"); + } + const auto [k_cache, v_cache] = + offgraph_cache_->layers[static_cast(layer)]; + if (query.sizes.size() != 4 || key.sizes.size() != 4 || + val.sizes.size() != 4 || query.sizes[0] != 1 || key.sizes[0] != 1 || + val.sizes[0] != 1) { + throw std::runtime_error( + "ggml: update_and_attend expects batch-1 [B,H,T,D] q/k/v"); + } + const int64_t tokens = query.sizes[2]; + if (key.sizes != val.sizes || key.sizes[2] != tokens || + key.sizes[3] != query.sizes[3] || key.sizes[1] <= 0 || + query.sizes[1] % key.sizes[1] != 0) { + throw std::runtime_error( + "ggml: update_and_attend q/k/v shapes are incompatible"); + } + if (k_cache->ne[0] != key.sizes[3] || k_cache->ne[2] != key.sizes[1]) { + throw std::runtime_error( + "ggml: update_and_attend layer " + std::to_string(layer) + + " K/V shape disagrees with the cache geometry"); + } + const int64_t keys = offgraph_cache_->keys; + ggml_context* ctx = instance_.context.get(); + if (!instance_.offgraph) { + ggml_tensor* rows = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, tokens); + ggml_set_name(rows, "kv_cache.rows"); + ggml_set_input(rows); + ggml_tensor* mask = new_tensor(ctx, GGML_TYPE_F16, {1, 1, tokens, keys}); + ggml_set_name(mask, "kv_cache.mask"); + ggml_set_input(mask); + instance_.offgraph = OffGraphUploads{ + .rows = instance_.uploads.size(), + .mask = instance_.uploads.size() + 1, + .tokens = tokens, + .keys = keys}; + instance_.uploads.push_back(Upload{.tensor = rows}); + instance_.uploads.push_back(Upload{.tensor = mask}); + } else if (instance_.offgraph->tokens != tokens) { + throw std::runtime_error( + "ggml: every update_and_attend in a method must see the same tokens"); + } + ggml_tensor* rows = instance_.uploads[instance_.offgraph->rows].tensor; + ggml_tensor* mask = instance_.uploads[instance_.offgraph->mask].tensor; + + const auto write = [&](ggml_tensor* cache, ggml_tensor* step) { + // set_rows converts F32 to either cache type, but not F16 to F32 on + // every backend. + ggml_tensor* source = + cache->type == GGML_TYPE_F16 && step->type == GGML_TYPE_F16 + ? skip_f16_cast(step) + : as_f32(step); + ggml_tensor* written = ggml_set_rows(ctx, cache, source, rows); + ggml_build_forward_expand(instance_.graph, written); + // Viewing the write orders the read after it. + return ggml_view_3d( + ctx, written, cache->ne[0], keys, cache->ne[2], cache->nb[1], + cache->nb[2], 0); + }; + ggml_tensor* k = write(k_cache, key.tensor); + ggml_tensor* v = write(v_cache, val.tensor); + + ggml_tensor* q = query.tensor; + if (k_cache->type == GGML_TYPE_F16) { + q = skip_f16_cast(q); + } + const auto scale = static_cast(float_arg(node, "scale", 5, 1.0)); + ggml_tensor* result = + ggml_flash_attn_ext(ctx, as_f32(q), k, v, mask, scale, 0.0f, 0.0f); + ggml_prec_set_acc(result, GGML_PREC_F32); + // flash_attn_ext returns [Dv, H, T, B]; torch order is [B, H, T, Dv]. + TensorRef out{ + ggml_permute(ctx, result, 0, 2, 1, 3), + {query.sizes[0], query.sizes[1], tokens, val.sizes[3]}}; + const std::vector ids = output_ids(node); + const ggml_type out_type = ids.empty() + ? GGML_TYPE_F32 + : compute_type(method_.graph.value(ids.front()).tensor_meta()); + if (out_type != GGML_TYPE_F32) { + out.tensor = ggml_cast(ctx, out.tensor, out_type); + set_single_output(node, out); + } else { + set_single_view(node, out); + } +} + } // namespace ptn::ggml_detail diff --git a/backends/native/runtime/engine/targets.bzl b/backends/native/runtime/engine/targets.bzl index a7ba9d20939..abb5c8e9488 100644 --- a/backends/native/runtime/engine/targets.bzl +++ b/backends/native/runtime/engine/targets.bzl @@ -49,6 +49,7 @@ def define_common_targets(): runtime.cxx_library( name = "ggml_engine_" + variant, srcs = [ + "ggml/GgmlCache.cpp", "ggml/GgmlEngine.cpp", "ggml/GgmlExecutor.cpp", "ggml/GgmlGraphBuilder.cpp", @@ -61,6 +62,7 @@ def define_common_targets(): "ggml/GgmlTypes.cpp", ], headers = [ + "ggml/GgmlCache.h", "ggml/GgmlExecutor.h", "ggml/GgmlGraphBuilder.h", "ggml/GgmlHostInts.h", @@ -81,6 +83,8 @@ def define_common_targets(): "//executorch/backends/native/runtime/graph:scalar_type", "//executorch/backends/native/runtime/graph:tensor_meta", "//executorch/backends/native/runtime/graph:value", + "//executorch/extension/llm/cache:kv_cache", + "//executorch/runtime/platform:platform", ggml_dep, ], labels = ggml_ci_labels(variant), diff --git a/backends/native/test/CMakeLists.txt b/backends/native/test/CMakeLists.txt index 73ea13f321c..55d42aea698 100644 --- a/backends/native/test/CMakeLists.txt +++ b/backends/native/test/CMakeLists.txt @@ -56,7 +56,7 @@ if(EXECUTORCH_BUILD_NATIVE_GGML) add_executable(native_ggml_engine_test runtime/ggml/test_ggml_engine.cpp) target_link_libraries( native_ggml_engine_test PRIVATE native_ggml_engine native_graph_schema - GTest::gtest_main + extension_llm_cache GTest::gtest_main ) add_dependencies(native_ggml_engine_test native_graph_schema_gen) add_test(NAME native_ggml_engine_test COMMAND native_ggml_engine_test) diff --git a/backends/native/test/runtime/ggml/BUCK b/backends/native/test/runtime/ggml/BUCK index 2b2add96f39..4146ae98be1 100644 --- a/backends/native/test/runtime/ggml/BUCK +++ b/backends/native/test/runtime/ggml/BUCK @@ -18,6 +18,7 @@ fbcode_target( "//executorch/backends/native/runtime/deserialize:package_test_data", "//executorch/backends/native/runtime/engine:engine", "//executorch/backends/native/runtime/engine:ggml_engine_cpu", + "//executorch/extension/llm/cache:kv_cache", ], **CPU_COMPAT ) diff --git a/backends/native/test/runtime/ggml/test_ggml_engine.cpp b/backends/native/test/runtime/ggml/test_ggml_engine.cpp index a7498fc9a07..147f5b9553d 100644 --- a/backends/native/test/runtime/ggml/test_ggml_engine.cpp +++ b/backends/native/test/runtime/ggml/test_ggml_engine.cpp @@ -29,6 +29,8 @@ #include #include #include +#include +#include #include #include @@ -234,7 +236,8 @@ std::vector make_linear_program( std::unique_ptr make_context( GgmlEngineHost& host, - const std::vector& program_bytes) { + const std::vector& program_bytes, + const EngineOptions& options = {}) { // One block_q4_K: d = 1.0 (fp16), dmin = 0, every sub-block scale 1 and // min 0, and nibbles 1 (low) / 2 (high), so half the weights are 1 and // half are 2. @@ -257,7 +260,8 @@ std::unique_ptr make_context( })))); auto program = std::make_shared( Program::load(program_bytes.data(), program_bytes.size())); - return host.create_ggml_context(std::move(program), std::move(package)); + return host.create_ggml_context( + std::move(program), std::move(package), options); } // A tensor (by name), an int, a float, an int list, or None. @@ -1329,6 +1333,386 @@ TEST(GgmlEngineTest, CustomSdpaMatchesReferenceWithoutMask) { } } +namespace kv = ::executorch::extension::llm::cache; + +constexpr int64_t kOffGraphHeads = 4; +constexpr int64_t kOffGraphKvHeads = 2; +constexpr int64_t kOffGraphDim = 8; +const double kOffGraphScale = 1.0 / std::sqrt(static_cast(kOffGraphDim)); + +// One kvcache::update_and_attend over q [1, 4, T, 8] and k, v [1, 2, T, 8] at +// positions [T, 1], addressing cache layer `layer`. With `gather_row`, the +// program returns the row of the output, viewed as [4 * T, 8], that a fifth +// input [1] names; the engine checks it on the host while building the graph. +std::vector make_update_and_attend_program( + int64_t tokens, + int64_t layer = 0, + bool gather_row = false) { + TestProgram program; + program.input( + "q", fbs::ScalarType::FLOAT, {1, kOffGraphHeads, tokens, kOffGraphDim}); + program.input( + "k", fbs::ScalarType::FLOAT, {1, kOffGraphKvHeads, tokens, kOffGraphDim}); + program.input( + "v", fbs::ScalarType::FLOAT, {1, kOffGraphKvHeads, tokens, kOffGraphDim}); + program.input("position", fbs::ScalarType::LONG, {tokens, 1}); + if (gather_row) { + program.input("row", fbs::ScalarType::LONG, {1}); + } + program.call( + "attention", + "torch.ops.kvcache.update_and_attend.default", + {{"q", std::string("q")}, + {"k", std::string("k")}, + {"v", std::string("v")}, + {"position", std::string("position")}, + {"layer_id", layer}, + {"scale", kOffGraphScale}}, + fbs::ScalarType::FLOAT, + {1, kOffGraphHeads, tokens, kOffGraphDim}); + if (!gather_row) { + return program.finish("attention"); + } + program.call( + "rows", + "torch.ops.aten.view_copy.default", + {{"self", std::string("attention")}, + {"size", std::vector{kOffGraphHeads * tokens, kOffGraphDim}}}, + fbs::ScalarType::FLOAT, + {kOffGraphHeads * tokens, kOffGraphDim}); + program.call( + "selected", + "torch.ops.aten.embedding.default", + {{"weight", std::string("rows")}, {"indices", std::string("row")}}, + fbs::ScalarType::FLOAT, + {1, kOffGraphDim}); + return program.finish("selected"); +} + +// A ggml cache built through the factory and published under a registry key, +// with the context option that names it. +struct InstalledCache { + std::shared_ptr cache; + std::unique_ptr guard; + EngineOptions options; + + kv::SequenceControl& control() { + return *cache->as(); + } +}; + +InstalledCache install_cache( + int capacity, + int initial_capacity, + ScalarType kv_dtype = ScalarType::Float, + int head_dim = static_cast(kOffGraphDim)) { + register_ggml_cache_builders(); + kv::CacheGeometry geometry; + geometry.layers.push_back(kv::LayerGeometry{ + .n_kv_heads = static_cast(kOffGraphKvHeads), .head_dim = head_dim}); + const kv::CacheConfig cfg{ + .capacity = capacity, + .kv_dtype = static_cast(kv_dtype), + .initial_capacity = initial_capacity}; + auto built = kv::CacheFactory::global().build( + kGgmlCacheBackendId, kv::kind::kSingle, geometry, cfg); + if (!built.ok()) { + throw std::runtime_error("failed to build the test cache"); + } + InstalledCache out; + out.cache = built.get(); + out.guard = std::make_unique(out.cache); + ::executorch::runtime::BackendOptions<1> options; + const char* key = nullptr; + if (out.guard->set_option(options) != ::executorch::runtime::Error::Ok || + options.get_option(kv::kCacheKeyOption, key) != + ::executorch::runtime::Error::Ok) { + throw std::runtime_error("failed to publish the test cache"); + } + out.options.emplace(kv::kCacheKeyOption, key); + return out; +} + +// What the cache should hold: K and V [Hkv, D] per position, written by each +// step, attended causally by naive softmax attention. +class OffGraphReference { + public: + // Runs `tokens` tokens from `start` through `executable`, with inputs that + // depend on the position and `seed`, and returns the largest difference + // from the reference. + float step( + EngineExecutable& executable, + int64_t start, + int64_t tokens, + int seed = 0) { + const int64_t kv_row = kOffGraphKvHeads * kOffGraphDim; + std::vector q(kOffGraphHeads * tokens * kOffGraphDim); + std::vector k(kOffGraphKvHeads * tokens * kOffGraphDim); + std::vector v(k.size()); + std::vector positions(tokens); + for (int64_t t = 0; t < tokens; ++t) { + const int64_t pos = start + t; + positions[t] = pos; + if (static_cast(k_.size()) <= pos) { + k_.resize(pos + 1, std::vector(kv_row)); + v_.resize(pos + 1, std::vector(kv_row)); + } + for (int64_t h = 0; h < kOffGraphHeads; ++h) { + for (int64_t d = 0; d < kOffGraphDim; ++d) { + q[(h * tokens + t) * kOffGraphDim + d] = static_cast( + std::sin(0.17 * pos + 0.9 * h - 0.21 * d + seed)); + } + } + for (int64_t h = 0; h < kOffGraphKvHeads; ++h) { + for (int64_t d = 0; d < kOffGraphDim; ++d) { + const auto key = static_cast( + std::sin(0.3 * pos + 0.7 * h + 0.13 * d + 2 * seed)); + const auto value = static_cast( + std::cos(0.2 * pos - 0.5 * h + 0.31 * d + seed) * (1 + pos % 3)); + k[(h * tokens + t) * kOffGraphDim + d] = key; + v[(h * tokens + t) * kOffGraphDim + d] = value; + k_[pos][h * kOffGraphDim + d] = key; + v_[pos][h * kOffGraphDim + d] = value; + } + } + } + executable.set_input(0, q.data(), q.size(), ScalarType::Float); + executable.set_input(1, k.data(), k.size(), ScalarType::Float); + executable.set_input(2, v.data(), v.size(), ScalarType::Float); + executable.set_input( + 3, positions.data(), positions.size(), ScalarType::Long); + executable.execute(); + std::vector out(q.size()); + executable.get_output(0, out.data(), out.size(), ScalarType::Float); + + float worst = 0.0f; + for (int64_t h = 0; h < kOffGraphHeads; ++h) { + const int64_t kv_head = h / (kOffGraphHeads / kOffGraphKvHeads); + for (int64_t t = 0; t < tokens; ++t) { + const float* query = &q[(h * tokens + t) * kOffGraphDim]; + const int64_t keys = start + t + 1; + std::vector weights(keys); + for (int64_t s = 0; s < keys; ++s) { + weights[s] = std::exp( + std::inner_product( + query, + query + kOffGraphDim, + &k_[s][kv_head * kOffGraphDim], + 0.0) * + kOffGraphScale); + } + const double total = + std::accumulate(weights.begin(), weights.end(), 0.0); + for (int64_t d = 0; d < kOffGraphDim; ++d) { + double sum = 0.0; + for (int64_t s = 0; s < keys; ++s) { + sum += weights[s] * v_[s][kv_head * kOffGraphDim + d]; + } + worst = std::max( + worst, + std::abs( + out[(h * tokens + t) * kOffGraphDim + d] - + static_cast(sum / total))); + } + } + } + return worst; + } + + private: + std::vector> k_; + std::vector> v_; +}; + +// A three-token prefill and decode steps in two contexts sharing one cache, +// across graph caching modes and both storage types. Storage starts at one +// row, so it regrows in the prefill and twice while decoding. +TEST(GgmlEngineTest, OffGraphAttentionMatchesReferenceAsTheCacheGrows) { + for (const size_t graphs : {size_t{0}, size_t{1}, size_t{4}}) { + for (const auto& [dtype, tol] : + {std::pair{ScalarType::Float, 1e-4f}, + std::pair{ScalarType::Half, 2e-2f}}) { + SCOPED_TRACE( + "graphs " + std::to_string(graphs) + ", " + + std::string(scalar_type_name(dtype))); + InstalledCache cache = install_cache(64, 1, dtype); + auto host = GgmlEngineHost::create({.max_cached_graphs = graphs}); + auto prefill_context = make_context( + *host, make_update_and_attend_program(3), cache.options); + auto decode_context = make_context( + *host, make_update_and_attend_program(1), cache.options); + auto prefill = prefill_context->compile("forward"); + auto decode = decode_context->compile("forward"); + OffGraphReference reference; + EXPECT_LT(reference.step(*prefill, 0, 3), tol); + for (int64_t pos = 3; pos < 12; ++pos) { + SCOPED_TRACE(pos); + EXPECT_LT(reference.step(*decode, pos, 1), tol); + } + } + } +} + +// After a clear, the old rows are masked out, including in an executable +// whose storage binding went stale when another regrew the cache. +TEST(GgmlEngineTest, OffGraphAttentionReplaysFromEmptyAfterClear) { + InstalledCache cache = install_cache(64, 2); + auto host = GgmlEngineHost::create(); + auto prefill_context = + make_context(*host, make_update_and_attend_program(3), cache.options); + auto decode_context = + make_context(*host, make_update_and_attend_program(1), cache.options); + auto prefill = prefill_context->compile("forward"); + auto decode = decode_context->compile("forward"); + { + OffGraphReference reference; + EXPECT_LT(reference.step(*prefill, 0, 3), 1e-4f); + for (int64_t pos = 3; pos < 10; ++pos) { + EXPECT_LT(reference.step(*decode, pos, 1), 1e-4f); + } + } + cache.control().clear(); + OffGraphReference reference; + EXPECT_LT(reference.step(*prefill, 0, 3, /*seed=*/5), 1e-4f); + EXPECT_LT(reference.step(*decode, 3, 1, /*seed=*/5), 1e-4f); +} + +// After a rewind, the next step overwrites the rewound rows and attends +// nothing past them. +TEST(GgmlEngineTest, OffGraphAttentionContinuesFromARewind) { + InstalledCache cache = install_cache(64, 4); + auto host = GgmlEngineHost::create(); + auto prefill_context = + make_context(*host, make_update_and_attend_program(3), cache.options); + auto decode_context = + make_context(*host, make_update_and_attend_program(1), cache.options); + auto prefill = prefill_context->compile("forward"); + auto decode = decode_context->compile("forward"); + OffGraphReference reference; + EXPECT_LT(reference.step(*prefill, 0, 3), 1e-4f); + for (int64_t pos = 3; pos < 8; ++pos) { + EXPECT_LT(reference.step(*decode, pos, 1), 1e-4f); + } + ASSERT_TRUE(cache.control().rewind(2)); + EXPECT_LT(reference.step(*decode, 2, 1, /*seed=*/9), 1e-4f); + EXPECT_LT(reference.step(*decode, 3, 1, /*seed=*/9), 1e-4f); +} + +// A step that starts before the sequence's length rewinds to its start, as a +// re-run of the same positions (warmup) does; one that starts past the length +// would attend rows nothing wrote, so it is rejected. +TEST(GgmlEngineTest, OffGraphStepsStartAtOrBeforeTheLength) { + InstalledCache cache = install_cache(64, 4); + auto host = GgmlEngineHost::create(); + auto prefill_context = + make_context(*host, make_update_and_attend_program(3), cache.options); + auto decode_context = + make_context(*host, make_update_and_attend_program(1), cache.options); + auto prefill = prefill_context->compile("forward"); + auto decode = decode_context->compile("forward"); + OffGraphReference reference; + EXPECT_LT(reference.step(*prefill, 0, 3), 1e-4f); + EXPECT_LT(reference.step(*prefill, 0, 3), 1e-4f); + // The length is 3 again. + EXPECT_THROW(reference.step(*decode, 4, 1), std::invalid_argument); + for (int64_t pos = 3; pos < 8; ++pos) { + EXPECT_LT(reference.step(*decode, pos, 1), 1e-4f); + } + // Back to position 4, so the length is 5. + EXPECT_LT(reference.step(*decode, 4, 1, /*seed=*/3), 1e-4f); + EXPECT_THROW(reference.step(*decode, 6, 1), std::invalid_argument); + EXPECT_LT(reference.step(*decode, 5, 1, /*seed=*/3), 1e-4f); +} + +// A step that fails before its graph runs, here on a row outside the output, +// leaves the sequence's length where it was, even if it started earlier. +TEST(GgmlEngineTest, FailedOffGraphStepLeavesTheSequence) { + InstalledCache cache = install_cache(64, 4); + auto host = GgmlEngineHost::create(); + auto prefill_context = + make_context(*host, make_update_and_attend_program(3), cache.options); + auto decode_context = + make_context(*host, make_update_and_attend_program(1), cache.options); + auto gather_context = make_context( + *host, make_update_and_attend_program(1, 0, true), cache.options); + auto prefill = prefill_context->compile("forward"); + auto decode = decode_context->compile("forward"); + auto gather = gather_context->compile("forward"); + OffGraphReference reference; + EXPECT_LT(reference.step(*prefill, 0, 3), 1e-4f); + for (int64_t pos = 3; pos < 6; ++pos) { + EXPECT_LT(reference.step(*decode, pos, 1), 1e-4f); + } + const std::vector q(kOffGraphHeads * kOffGraphDim); + const std::vector kv_step(kOffGraphKvHeads * kOffGraphDim); + const int64_t position = 2; + const int64_t row = kOffGraphHeads; + gather->set_input(0, q.data(), q.size(), ScalarType::Float); + gather->set_input(1, kv_step.data(), kv_step.size(), ScalarType::Float); + gather->set_input(2, kv_step.data(), kv_step.size(), ScalarType::Float); + gather->set_input(3, &position, 1, ScalarType::Long); + gather->set_input(4, &row, 1, ScalarType::Long); + EXPECT_THROW(gather->execute(), std::out_of_range); + EXPECT_LT(reference.step(*decode, 6, 1), 1e-4f); +} + +TEST(GgmlEngineTest, RejectsOffGraphStepsPastCapacityOrNotConsecutive) { + InstalledCache cache = install_cache(4, 4); + auto host = GgmlEngineHost::create(); + auto prefill_context = + make_context(*host, make_update_and_attend_program(3), cache.options); + auto decode_context = + make_context(*host, make_update_and_attend_program(1), cache.options); + auto prefill = prefill_context->compile("forward"); + auto decode = decode_context->compile("forward"); + OffGraphReference reference; + EXPECT_THROW(reference.step(*decode, 1, 1), std::invalid_argument); + EXPECT_LT(reference.step(*prefill, 0, 3), 1e-4f); + EXPECT_LT(reference.step(*decode, 3, 1), 1e-4f); + EXPECT_THROW(reference.step(*decode, 4, 1), std::out_of_range); + EXPECT_THROW(reference.step(*prefill, 2, 3), std::out_of_range); + + const std::vector q(kOffGraphHeads * 3 * kOffGraphDim); + const std::vector kv_step(kOffGraphKvHeads * 3 * kOffGraphDim); + const std::vector positions = {0, 2, 3}; + prefill->set_input(0, q.data(), q.size(), ScalarType::Float); + prefill->set_input(1, kv_step.data(), kv_step.size(), ScalarType::Float); + prefill->set_input(2, kv_step.data(), kv_step.size(), ScalarType::Float); + prefill->set_input(3, positions.data(), positions.size(), ScalarType::Long); + EXPECT_THROW(prefill->execute(), std::invalid_argument); +} + +TEST(GgmlEngineTest, OffGraphAttentionNeedsAMatchingCache) { + auto host = GgmlEngineHost::create(); + // No cache option. + EXPECT_THROW( + make_context(*host, make_update_and_attend_program(1)) + ->compile("forward"), + std::runtime_error); + // A key nothing is installed under, and an option the engine does not know. + EXPECT_THROW( + make_context( + *host, + make_update_and_attend_program(1), + {{std::string(kv::kCacheKeyOption), "no-such-cache"}}), + std::invalid_argument); + EXPECT_THROW( + make_context( + *host, make_update_and_attend_program(1), {{"unknown", "value"}}), + std::invalid_argument); + // A layer the cache does not have, and K/V shaped unlike the cache. + InstalledCache cache = install_cache(8, 8); + EXPECT_THROW( + make_context(*host, make_update_and_attend_program(1, 1), cache.options) + ->compile("forward"), + std::runtime_error); + InstalledCache narrow = install_cache(8, 8, ScalarType::Float, 4); + EXPECT_THROW( + make_context(*host, make_update_and_attend_program(1), narrow.options) + ->compile("forward"), + std::runtime_error); +} + TEST(GgmlEngineTest, RejectsQuantizedGraphInputsAndOutputs) { TestProgram as_input; as_input.input("x", fbs::ScalarType::FLOAT, {1, 256});