diff --git a/backends/native/runtime/engine/ggml/GgmlCache.cpp b/backends/native/runtime/engine/ggml/GgmlCache.cpp index 0c79e9d3fc8..880e089fd4d 100644 --- a/backends/native/runtime/engine/ggml/GgmlCache.cpp +++ b/backends/native/runtime/engine/ggml/GgmlCache.cpp @@ -12,6 +12,7 @@ #include #include #include +#include #include #include @@ -70,16 +71,30 @@ std::vector ring_rows( } // namespace +GgmlKvPool::GgmlKvPool( + const cache::CacheGeometry& geometry, + int kv_dtype, + int initial_rows, + int capacity, + std::vector ring_rows) + : geometry_(geometry), + ring_rows_(std::move(ring_rows)), + type_(kv_type(kv_dtype)), + initial_rows_(std::clamp(initial_rows, 1, capacity)), + capacity_(capacity) {} + GgmlSequenceCache::GgmlSequenceCache( const cache::CacheGeometry& geometry, const cache::CacheConfig& cfg) : SequenceCache(checked(geometry, cfg), cfg), - geometry_(geometry), - ring_rows_(ring_rows(geometry, cfg)), - 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) { + pool_( + geometry, + cfg.kv_dtype, + cfg.initial_capacity, + cfg.capacity, + ring_rows(geometry, cfg)) {} + +void GgmlKvPool::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_) { @@ -88,10 +103,10 @@ void GgmlSequenceCache::reserve(ggml_backend_t backend, int rows) { ggml_backend_buft_name(buffer_type_) + ", not " + ggml_backend_buft_name(buffer_type)); } - if (rows > capacity()) { + if (rows > capacity_) { throw std::out_of_range( "ggml cache: " + std::to_string(rows) + - " rows exceed the capacity of " + std::to_string(capacity())); + " rows exceed the capacity of " + std::to_string(capacity_)); } if (storage_ && rows <= storage_->rows) { return; @@ -101,7 +116,7 @@ void GgmlSequenceCache::reserve(ggml_backend_t backend, int rows) { target *= 2; } auto next = std::make_shared(); - next->rows = static_cast(std::min(target, capacity())); + next->rows = static_cast(std::min(target, capacity_)); ggml_init_params params{ .mem_size = ggml_tensor_overhead() * (2 * geometry_.layers.size() + 1), @@ -114,7 +129,7 @@ void GgmlSequenceCache::reserve(ggml_backend_t backend, int rows) { } for (size_t i = 0; i < geometry_.layers.size(); ++i) { const cache::LayerGeometry& layer = geometry_.layers[i]; - Layer out{.window = layer.policy.window}; + Layer out{.window = layer.policy.window, .ring = ring_rows_[i] > 0}; for (ggml_tensor** t : {&out.k, &out.v}) { *t = ggml_new_tensor_3d( next->context.get(), diff --git a/backends/native/runtime/engine/ggml/GgmlCache.h b/backends/native/runtime/engine/ggml/GgmlCache.h index dc44c39f3dc..4477b1e318a 100644 --- a/backends/native/runtime/engine/ggml/GgmlCache.h +++ b/backends/native/runtime/engine/ggml/GgmlCache.h @@ -19,25 +19,20 @@ 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]). A flat -// layer's row r holds logical position r; its storage starts at the config's -// initial_capacity rows and doubles toward capacity as steps need more. A ring -// layer holds its window + max_write - 1 rows from the start, and the -// SequencePlanner says which rows a step writes and reads. Like SequenceCache, -// it is not thread-safe: executables bound to one cache must be serialized. -class GgmlSequenceCache final - : public executorch::extension::llm::cache::SequenceCache { +// K/V storage per cache 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]). A growing layer starts at `initial_rows` +// and doubles toward `capacity` as steps need more; a ring layer holds its +// fixed rows from the start. Not thread-safe. +class GgmlKvPool { public: - static constexpr const char* kFaceName = "ptn.ggml.SequenceCache"; - struct Layer { ggml_tensor* k = nullptr; ggml_tensor* v = nullptr; - // Sliding window of a ring layer; 0 for a flat one. + // The layer's sliding window, 0 for none; layers sharing one share a mask. int window = 0; + // A ring layer attends all its fixed rows; a growing one, a prefix. + bool ring = false; }; // One allocation of K/V storage. A regrow moves the rows into a new one; @@ -46,21 +41,23 @@ class GgmlSequenceCache final ggml_context_ptr context; ggml_backend_buffer_ptr buffer; std::vector layers; - // Rows of every flat layer. + // Rows of every growing layer. int rows = 0; }; - // Throws on geometry or config the ggml engine cannot serve. - GgmlSequenceCache( + // `ring_rows` holds each layer's fixed rows, or 0 for a growing layer. + // Throws for a kv_dtype other than Float or Half. + GgmlKvPool( const executorch::extension::llm::cache::CacheGeometry& geometry, - const executorch::extension::llm::cache::CacheConfig& cfg); - - GgmlSequenceCache(const GgmlSequenceCache&) = delete; - GgmlSequenceCache& operator=(const GgmlSequenceCache&) = delete; + int kv_dtype, + int initial_rows, + int capacity, + std::vector ring_rows); - // Ensures flat layers 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. + // Ensures growing layers 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 @@ -69,17 +66,44 @@ class GgmlSequenceCache final return storage_; } - protected: - void* face(executorch::extension::llm::cache::FaceId id) override; - private: executorch::extension::llm::cache::CacheGeometry geometry_; - // Per layer: the fixed rows of a ring layer, 0 for a flat one. std::vector ring_rows_; ggml_type type_; int initial_rows_; + int capacity_; ggml_backend_buffer_type_t buffer_type_ = nullptr; std::shared_ptr storage_; }; +// The off-graph KV cache of one sequence for the ggml engine: the neutral +// SequenceCache controller over a GgmlKvPool. A flat layer's row r holds +// logical position r; a ring layer holds its window + max_write - 1 rows, and +// the SequencePlanner says which rows a step writes and reads. 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"; + + // 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; + + GgmlKvPool& pool() { + return pool_; + } + + protected: + void* face(executorch::extension::llm::cache::FaceId id) override; + + private: + GgmlKvPool pool_; +}; + } // namespace ptn diff --git a/backends/native/runtime/engine/ggml/GgmlExecutor.cpp b/backends/native/runtime/engine/ggml/GgmlExecutor.cpp index 22014de402d..76d1387e24e 100644 --- a/backends/native/runtime/engine/ggml/GgmlExecutor.cpp +++ b/backends/native/runtime/engine/ggml/GgmlExecutor.cpp @@ -1293,19 +1293,22 @@ class GgmlExecutable final : public EngineExecutable { // 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_) { + kv_cache_->pool().reserve(primary_backend_, static_cast(rows)); + if (kv_cache_->pool().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_storage_ = kv_cache_->pool().storage(); offgraph_cache_.layers.clear(); - for (const GgmlSequenceCache::Layer& layer : offgraph_storage_->layers) { + for (const GgmlKvPool::Layer& layer : offgraph_storage_->layers) { offgraph_cache_.layers.push_back( - {.k = layer.k, .v = layer.v, .window = layer.window}); + {.k = layer.k, + .v = layer.v, + .window = layer.window, + .ring = layer.ring}); } } @@ -1509,7 +1512,7 @@ class GgmlExecutable final : public EngineExecutable { 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_; + std::shared_ptr offgraph_storage_; // The step execute() is running, checked once. std::optional offgraph_step_; std::vector inputs_; diff --git a/backends/native/runtime/engine/ggml/GgmlGraphBuilder.h b/backends/native/runtime/engine/ggml/GgmlGraphBuilder.h index 84ba6df1dad..6c09cbc1809 100644 --- a/backends/native/runtime/engine/ggml/GgmlGraphBuilder.h +++ b/backends/native/runtime/engine/ggml/GgmlGraphBuilder.h @@ -94,12 +94,13 @@ struct OffGraphCacheBinding { // K and V storage, ggml [head_dim, rows, n_kv_heads]. ggml_tensor* k = nullptr; ggml_tensor* v = nullptr; - // Sliding window of a ring layer, which attends all its rows; 0 for a - // flat layer. + // The layer's sliding window, 0 for none; layers sharing one share a mask. int window = 0; + // A ring layer attends all its rows; any other, the first `keys`. + bool ring = false; }; std::vector layers; - // Rows a flat layer's attention reads, a prefix of its storage. + // Rows a non-ring layer's attention reads, a prefix of its storage. int64_t keys = 0; }; diff --git a/backends/native/runtime/engine/ggml/GgmlLowerAttention.cpp b/backends/native/runtime/engine/ggml/GgmlLowerAttention.cpp index 4f0e00210eb..a169e421150 100644 --- a/backends/native/runtime/engine/ggml/GgmlLowerAttention.cpp +++ b/backends/native/runtime/engine/ggml/GgmlLowerAttention.cpp @@ -398,8 +398,8 @@ void GraphBuilder::lower_sdpa(const Node& node) { } // Writes this step's K/V into the cache rows the executor uploads, then -// attends the layer's rows under the uploaded mask: a flat layer's prefix, or -// all of a ring layer. Q/K/V are [B,H,T,D]. +// attends the layer's rows under the uploaded mask: a prefix, or all of a +// ring layer. 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( @@ -417,7 +417,7 @@ void GraphBuilder::lower_update_and_attend(const Node& node) { "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, window] = + const auto [k_cache, v_cache, window, ring] = 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 || @@ -437,7 +437,7 @@ void GraphBuilder::lower_update_and_attend(const Node& node) { "ggml: update_and_attend layer " + std::to_string(layer) + " K/V shape disagrees with the cache geometry"); } - const int64_t keys = window > 0 ? k_cache->ne[1] : offgraph_cache_->keys; + const int64_t keys = ring ? k_cache->ne[1] : offgraph_cache_->keys; ggml_context* ctx = instance_.context.get(); if (!instance_.offgraph) { instance_.offgraph = OffGraphUploads{.tokens = tokens};