From 0d8ca8fed52a0f98e787cab83ba161c297ec7720 Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Thu, 12 Nov 2026 11:09:05 -0800 Subject: [PATCH] Serve many sequences from one ggml engine KV cache The ggml engine's off-graph cache held one sequence, so a batching runtime could not put several sessions in one forward, clone a prefix, or share a pool between sessions. GgmlCellCache is the neutral CellCache over a GgmlKvPool whose rows are its cells, registered for kind::kBatched and kind::kBatchedCell; it offers BatchControl to the control plane. A context accepts either cache. With a cell cache, the executor places the forward's tokens at their positions before building the graph (the control plane declared their sequences), grows storage to the step's read extent, and fills each window's cache rows and mask from the cache's step: rows are the step's cells, and the mask keeps each query to its own sequence's earlier cells within its window. A windowed layer keeps every cell, so a rewind or clone never needs history the cache dropped. A forward that fails after placement leaves its cells claimed, as a failed MLX forward does. --- .../native/runtime/engine/ggml/GgmlCache.cpp | 20 + .../native/runtime/engine/ggml/GgmlCache.h | 48 ++ .../native/runtime/engine/ggml/GgmlEngine.cpp | 39 +- .../native/runtime/engine/ggml/GgmlEngine.h | 14 +- .../runtime/engine/ggml/GgmlExecutor.cpp | 147 ++++-- .../native/runtime/engine/ggml/GgmlExecutor.h | 5 +- .../test/runtime/ggml/test_ggml_engine.cpp | 444 +++++++++++++++++- extension/llm/cache/cell_cache.cpp | 21 + extension/llm/cache/cell_cache.h | 10 + extension/llm/cache/test/cache_test.cpp | 28 ++ 10 files changed, 719 insertions(+), 57 deletions(-) diff --git a/backends/native/runtime/engine/ggml/GgmlCache.cpp b/backends/native/runtime/engine/ggml/GgmlCache.cpp index 880e089fd4d..2bd34cc607c 100644 --- a/backends/native/runtime/engine/ggml/GgmlCache.cpp +++ b/backends/native/runtime/engine/ggml/GgmlCache.cpp @@ -182,6 +182,22 @@ void* GgmlSequenceCache::face(cache::FaceId id) { GgmlSequenceCache>(this, id); } +GgmlCellCache::GgmlCellCache( + const cache::CacheGeometry& geometry, + const cache::CacheConfig& cfg) + : CellCache(checked(geometry, cfg), cfg), + pool_( + geometry, + cfg.kv_dtype, + cfg.initial_capacity, + cfg.capacity, + std::vector(geometry.layers.size(), 0)) {} + +void* GgmlCellCache::face(cache::FaceId id) { + return cache:: + expose(this, id); +} + namespace { // CacheFactory::build reports a null cache as an error, so a cache the engine @@ -205,6 +221,10 @@ void register_ggml_cache_builders() { // Registering twice is harmless; the first builder stays. (void)cache::CacheFactory::global().register_builder( kGgmlCacheBackendId, cache::kind::kSingle, builder()); + for (const char* kind : {cache::kind::kBatched, cache::kind::kBatchedCell}) { + (void)cache::CacheFactory::global().register_builder( + kGgmlCacheBackendId, kind, builder()); + } } } // namespace ptn diff --git a/backends/native/runtime/engine/ggml/GgmlCache.h b/backends/native/runtime/engine/ggml/GgmlCache.h index 4477b1e318a..6372259e1c5 100644 --- a/backends/native/runtime/engine/ggml/GgmlCache.h +++ b/backends/native/runtime/engine/ggml/GgmlCache.h @@ -15,6 +15,7 @@ #include #include +#include #include namespace ptn { @@ -106,4 +107,51 @@ class GgmlSequenceCache final GgmlKvPool pool_; }; +// The off-graph KV cache of many sequences for the ggml engine: the neutral +// CellCache over a GgmlKvPool whose rows are its cells, every layer growing +// toward the pool's capacity. A step's cells and masks come from place_step; a +// windowed layer keeps every cell and its mask narrows to the window. Like +// CellCache, it is not thread-safe: executables bound to one cache must be +// serialized. +class GgmlCellCache final + : public executorch::extension::llm::cache::CellCache { + public: + static constexpr const char* kFaceName = "ptn.ggml.CellCache"; + + // Throws on geometry or config the ggml engine cannot serve. + GgmlCellCache( + const executorch::extension::llm::cache::CacheGeometry& geometry, + const executorch::extension::llm::cache::CacheConfig& cfg); + + GgmlCellCache(const GgmlCellCache&) = delete; + GgmlCellCache& operator=(const GgmlCellCache&) = delete; + + GgmlKvPool& pool() { + return pool_; + } + + protected: + void* face(executorch::extension::llm::cache::FaceId id) override; + + private: + GgmlKvPool pool_; +}; + +// The off-graph cache an executable binds: its storage and its layout, one of +// a single sequence or cells. +struct GgmlKvCache { + GgmlSequenceCache* sequence = nullptr; + GgmlCellCache* cells = nullptr; + + explicit operator bool() const { + return sequence != nullptr || cells != nullptr; + } + GgmlKvPool& pool() const { + return sequence != nullptr ? sequence->pool() : cells->pool(); + } + int capacity() const { + return sequence != nullptr ? sequence->capacity() : cells->capacity(); + } +}; + } // namespace ptn diff --git a/backends/native/runtime/engine/ggml/GgmlEngine.cpp b/backends/native/runtime/engine/ggml/GgmlEngine.cpp index 70429eb8c25..1617d4ccc9f 100644 --- a/backends/native/runtime/engine/ggml/GgmlEngine.cpp +++ b/backends/native/runtime/engine/ggml/GgmlEngine.cpp @@ -95,40 +95,47 @@ std::unique_ptr GgmlEngineHost::create_ggml_context( std::shared_ptr package, const EngineOptions& options) { namespace cache = executorch::extension::llm::cache; - std::shared_ptr kv_cache; + std::shared_ptr installed; + GgmlKvCache 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); + 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) { + kv_cache = { + .sequence = installed->as(), + .cells = installed->as()}; + if (!kv_cache) { throw std::invalid_argument( impl_->name + ": KV cache '" + value + "' was not built for this engine (CacheFactory backend '" + - kGgmlCacheBackendId + "', kind 'single')"); + kGgmlCacheBackendId + "')"); } - kv_cache = std::shared_ptr(installed, ggml_cache); } return std::unique_ptr(new GgmlEngineContext( std::move(program), std::move(package), impl_->options, - std::move(kv_cache))); + std::move(installed), + kv_cache)); } struct GgmlEngineContext::Impl { - Impl(GgmlEngineOptions value, std::shared_ptr cache) - : options(std::move(value)), kv_cache(std::move(cache)) {} + Impl( + GgmlEngineOptions value, + std::shared_ptr owner, + GgmlKvCache cache) + : options(std::move(value)), cache(std::move(owner)), kv_cache(cache) {} GgmlEngineOptions options; - std::shared_ptr kv_cache; + // Keeps the cache kv_cache points into alive. + std::shared_ptr cache; + GgmlKvCache kv_cache; size_t materialized_constant_bytes = 0; size_t heap_copied_constant_bytes = 0; GgmlMutableKeyOwners mutable_key_owners; @@ -138,9 +145,13 @@ GgmlEngineContext::GgmlEngineContext( std::shared_ptr program, std::shared_ptr package, GgmlEngineOptions options, - std::shared_ptr cache) + std::shared_ptr cache, + GgmlKvCache kv_cache) : EngineContext(std::move(program), std::move(package)), - impl_(std::make_unique(std::move(options), std::move(cache))) {} + impl_(std::make_unique( + std::move(options), + std::move(cache), + kv_cache)) {} GgmlEngineContext::~GgmlEngineContext() = default; @@ -153,7 +164,7 @@ std::unique_ptr GgmlEngineContext::compile_method( &impl_->materialized_constant_bytes, &impl_->heap_copied_constant_bytes, impl_->mutable_key_owners, - impl_->kv_cache.get()); + impl_->kv_cache); } 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 0269baa79f5..c8407b01724 100644 --- a/backends/native/runtime/engine/ggml/GgmlEngine.h +++ b/backends/native/runtime/engine/ggml/GgmlEngine.h @@ -12,12 +12,17 @@ #include +namespace executorch::extension::llm::cache { +class Cache; +} // namespace executorch::extension::llm::cache + 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. +// built under this id: kind::kSingle for one sequence, or kind::kBatched / +// kind::kBatchedCell for many sequences over shared cells. inline constexpr char kGgmlCacheBackendId[] = "NativeGgml"; // Registers the ggml cache builders with CacheFactory::global(). Idempotent. @@ -52,7 +57,7 @@ struct GgmlEngineOptions { }; class GgmlEngineContext; -class GgmlSequenceCache; +struct GgmlKvCache; class GgmlEngineHost final : public EngineHost { public: @@ -68,7 +73,7 @@ class GgmlEngineHost final : public EngineHost { std::shared_ptr package) override; // Accepts kCacheKeyOption, resolved through CacheRegistry::global(). Contexts - // bound to one cache share its sequence: the caller must serialize them. + // bound to one cache share its sequences: the caller must serialize them. std::unique_ptr create_context_with_options( std::shared_ptr program, std::shared_ptr package, @@ -105,7 +110,8 @@ class GgmlEngineContext final : public EngineContext { std::shared_ptr program, std::shared_ptr package, GgmlEngineOptions options, - std::shared_ptr cache); + std::shared_ptr cache, + GgmlKvCache kv_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 76d1387e24e..28c474e88b2 100644 --- a/backends/native/runtime/engine/ggml/GgmlExecutor.cpp +++ b/backends/native/runtime/engine/ggml/GgmlExecutor.cpp @@ -50,6 +50,7 @@ namespace ptn { namespace { using namespace ggml_detail; +namespace cache = executorch::extension::llm::cache; using executorch::extension::llm::cache::Run; using executorch::extension::llm::cache::SeqStepPlan; @@ -166,7 +167,7 @@ class GgmlExecutable final : public EngineExecutable { size_t* materialized_constant_bytes, size_t* heap_copied_constant_bytes, GgmlMutableKeyOwners& mutable_key_owners, - GgmlSequenceCache* kv_cache) + GgmlKvCache kv_cache) : mutable_key_claim_(mutable_key_owners, method), method_(method), package_(package), @@ -267,13 +268,15 @@ class GgmlExecutable final : public EngineExecutable { offgraph_step_ = checked_offgraph_step(); } - // A step that fails before its graph runs leaves the cache unchanged; one - // that fails while running may have written rows from its start. + // A cell step ends with its forward, giving back the cells it claimed if + // the forward fails. A sequence step that fails before its graph runs + // leaves the sequence 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); + bind_offgraph(offgraph_step_->rows); } if (options_.max_cached_graphs == 0) { ggml_backend_sched_reset(scheduler_.get()); @@ -301,17 +304,23 @@ class GgmlExecutable final : public EngineExecutable { } } } catch (...) { - if (ran && valid(offgraph_position_) && - offgraph_step_->start < kv_cache_->length()) { - kv_cache_->rewind(static_cast(offgraph_step_->start)); + if (valid(offgraph_position_) && kv_cache_.cells != nullptr) { + kv_cache_.cells->abandon_step(); + } else if ( + ran && valid(offgraph_position_) && + offgraph_step_->positions.front() < kv_cache_.sequence->length()) { + kv_cache_.sequence->rewind(offgraph_step_->positions.front()); } throw; } if (step) { - if (offgraph_step_->start < kv_cache_->length()) { - kv_cache_->rewind(static_cast(offgraph_step_->start)); + if (offgraph_step_->positions.front() < kv_cache_.sequence->length()) { + kv_cache_.sequence->rewind(offgraph_step_->positions.front()); } - kv_cache_->commit(*step); + kv_cache_.sequence->commit(*step); + } + if (valid(offgraph_position_) && kv_cache_.cells != nullptr) { + kv_cache_.cells->end_step(); } executed_ = true; } @@ -1246,8 +1255,8 @@ class GgmlExecutable final : public EngineExecutable { : 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; + // The cache rows the step reads, bucketed like masks. + const int64_t needed = offgraph_step_->rows; out.push_back(std::min( offgraph_storage_->rows, (needed + kKeyExtentBucket - 1) / kKeyExtentBucket * @@ -1280,7 +1289,7 @@ class GgmlExecutable final : public EngineExecutable { if (!valid(offgraph_position_)) { return; } - if (kv_cache_ == nullptr) { + if (!kv_cache_) { throw std::runtime_error( "ggml: method '" + method_.name + "' calls kvcache.update_and_attend but its context has no KV " @@ -1293,15 +1302,15 @@ 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_->pool().reserve(primary_backend_, static_cast(rows)); - if (kv_cache_->pool().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_->pool().storage(); + offgraph_storage_ = kv_cache_.pool().storage(); offgraph_cache_.layers.clear(); for (const GgmlKvPool::Layer& layer : offgraph_storage_->layers) { offgraph_cache_.layers.push_back( @@ -1313,14 +1322,13 @@ class GgmlExecutable final : public EngineExecutable { } struct OffGraphStep { - int64_t start; - int64_t tokens; + // Per token, its logical position. + std::vector positions; + // Cache rows the step reads, which storage must hold. + int64_t rows; }; - // 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. + // Reads the step's positions on the host and checks them against the cache. OffGraphStep checked_offgraph_step() { HostIntEvaluator host = host_evaluator(); const HostInts* values = host.tensor(offgraph_position_); @@ -1329,7 +1337,38 @@ class GgmlExecutable final : public EngineExecutable { "ggml: update_and_attend positions must be computable on the host " "from graph inputs"); } - const std::vector& positions = values->data; + OffGraphStep step; + for (const int64_t position : values->data) { + if (position < 0 || position >= kv_cache_.capacity()) { + throw std::out_of_range( + "ggml: KV cache position " + std::to_string(position) + + " is outside the cache capacity of " + + std::to_string(kv_cache_.capacity())); + } + step.positions.push_back(static_cast(position)); + } + if (kv_cache_.cells != nullptr) { + // Placing the step claims its cells for every layer of the forward. + const cache::CellStep* placed = kv_cache_.cells->place_step( + 0, step.positions.data(), static_cast(step.positions.size())); + if (placed == nullptr) { + throw std::invalid_argument( + "ggml: the cell cache refused the step: declare it first, with " + "positions continuing each sequence and room in the pool"); + } + step.rows = placed->read_len; + return step; + } + step.rows = checked_sequence_step(step.positions); + return step; + } + + // Checks a single-sequence step's positions are consecutive, fit the cache, + // and start no later than the sequence's length, at a position the sequence + // can rewind to, and returns the rows it reads. A step that starts earlier + // rewinds the sequence to its start once it has run. + int64_t checked_sequence_step(const std::vector& positions) { + GgmlSequenceCache& sequence = *kv_cache_.sequence; const int64_t start = positions.front(); for (size_t i = 0; i < positions.size(); ++i) { if (positions[i] != start + static_cast(i)) { @@ -1338,25 +1377,25 @@ class GgmlExecutable final : public EngineExecutable { } } const auto tokens = static_cast(positions.size()); - if (start < 0 || start + tokens > kv_cache_->capacity()) { + if (start + tokens > sequence.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())); + std::to_string(sequence.capacity())); } - if (start > kv_cache_->length()) { + if (start > sequence.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())); + std::to_string(sequence.length())); } - if (start < kv_cache_->length() && - !kv_cache_->can_rewind(static_cast(start))) { + if (start < sequence.length() && + !sequence.can_rewind(static_cast(start))) { throw std::runtime_error( "ggml: the KV cache cannot rewind to position " + std::to_string(start)); } - return {start, tokens}; + return start + tokens; } // Plans the step on the cache and fills each window's cache rows and mask. @@ -1366,13 +1405,18 @@ class GgmlExecutable final : public EngineExecutable { return std::nullopt; } const OffGraphUploads& slots = *instance.offgraph; - if (offgraph_step_->tokens != slots.tokens) { + if (static_cast(offgraph_step_->positions.size()) != + slots.tokens) { throw std::runtime_error( "ggml: update_and_attend needs one position per token"); } - const int64_t start = offgraph_step_->start; + if (kv_cache_.cells != nullptr) { + prepare_cell_step(instance); + return std::nullopt; + } + const int64_t start = offgraph_step_->positions.front(); const auto plan_layer = [&](int layer) { - const std::optional plan = kv_cache_->plan( + const std::optional plan = kv_cache_.sequence->plan( layer, static_cast(start), static_cast(slots.tokens)); if (!plan) { throw std::runtime_error( @@ -1436,6 +1480,39 @@ class GgmlExecutable final : public EngineExecutable { return committed; } + // Fills each window's cache rows and mask from the placed cell step, which + // every layer of the forward shares. + void prepare_cell_step(GraphInstance& instance) { + const OffGraphUploads& slots = *instance.offgraph; + 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 (const OffGraphUploads::Group& group : slots.groups) { + const cache::CellStep* step = kv_cache_.cells->place_step( + group.layer, + offgraph_step_->positions.data(), + static_cast(slots.tokens)); + if (step == nullptr || step->read_len > group.keys || + static_cast(step->cells.size()) != slots.tokens) { + throw std::logic_error("ggml: the cell step does not fit the graph"); + } + std::vector& rows = instance.uploads[group.rows].bytes; + rows.resize(step->cells.size() * sizeof(int32_t)); + std::memcpy(rows.data(), step->cells.data(), rows.size()); + std::vector& mask = instance.uploads[group.mask].bytes; + mask.resize( + static_cast(slots.tokens * group.keys) * sizeof(ggml_fp16_t)); + auto* bias = reinterpret_cast(mask.data()); + for (int64_t i = 0; i < slots.tokens; ++i) { + const uint8_t* bits = step->mask_bits.data() + i * step->read_len; + for (int64_t j = 0; j < group.keys; ++j) { + bias[i * group.keys + j] = + j < step->read_len && bits[j] != 0 ? attend : masked; + } + } + } + } + // 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); @@ -1505,7 +1582,7 @@ class GgmlExecutable final : public EngineExecutable { const Method& method_; const Package& package_; GgmlEngineOptions options_; - GgmlSequenceCache* kv_cache_; + GgmlKvCache 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. @@ -1607,7 +1684,7 @@ std::unique_ptr make_ggml_executable( size_t* materialized_constant_bytes, size_t* heap_copied_constant_bytes, GgmlMutableKeyOwners& mutable_key_owners, - GgmlSequenceCache* kv_cache) { + GgmlKvCache kv_cache) { return std::make_unique( method, package, diff --git a/backends/native/runtime/engine/ggml/GgmlExecutor.h b/backends/native/runtime/engine/ggml/GgmlExecutor.h index 6b4505093fc..c13f945d4bd 100644 --- a/backends/native/runtime/engine/ggml/GgmlExecutor.h +++ b/backends/native/runtime/engine/ggml/GgmlExecutor.h @@ -14,6 +14,7 @@ #include #include +#include #include namespace ptn { @@ -31,7 +32,7 @@ using GgmlMutableKeyOwners = std::unordered_map; // The executable holds the method's mutable buffer keys in // `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. +// method with such calls and no cache does not compile. std::unique_ptr make_ggml_executable( const Method& method, const Package& package, @@ -39,6 +40,6 @@ std::unique_ptr make_ggml_executable( size_t* materialized_constant_bytes, size_t* heap_copied_constant_bytes, GgmlMutableKeyOwners& mutable_key_owners, - GgmlSequenceCache* kv_cache = nullptr); + GgmlKvCache kv_cache = {}); } // namespace ptn diff --git a/backends/native/test/runtime/ggml/test_ggml_engine.cpp b/backends/native/test/runtime/ggml/test_ggml_engine.cpp index aa234878f4f..38be8cf5d7e 100644 --- a/backends/native/test/runtime/ggml/test_ggml_engine.cpp +++ b/backends/native/test/runtime/ggml/test_ggml_engine.cpp @@ -1427,6 +1427,49 @@ std::vector make_update_and_attend_program( return program.finish("selected"); } +// make_update_and_attend_program's attention once per cache layer in +// [0, layers), over the same q, k and v, summed. +std::vector make_layered_attention_program( + int64_t tokens, + int64_t layers) { + 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}); + std::string sum; + for (int64_t layer = 0; layer < layers; ++layer) { + const std::string attention = "attention" + std::to_string(layer); + 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 (sum.empty()) { + sum = attention; + continue; + } + const std::string next = "sum" + std::to_string(layer); + program.call( + next, + "torch.ops.aten.add.Tensor", + {{"self", sum}, {"other", attention}}, + fbs::ScalarType::FLOAT, + {1, kOffGraphHeads, tokens, kOffGraphDim}); + sum = next; + } + return program.finish(sum); +} + // A ggml cache built through the factory and published under a registry key, // with the context option that names it. struct InstalledCache { @@ -1437,6 +1480,9 @@ struct InstalledCache { kv::SequenceControl& control() { return *cache->as(); } + kv::BatchControl& batch() { + return *cache->as(); + } }; // One layer per entry of `windows`: flat for 0, else a ring of that window. @@ -1446,7 +1492,8 @@ InstalledCache install_cache( ScalarType kv_dtype = ScalarType::Float, int head_dim = static_cast(kOffGraphDim), const std::vector& windows = {0}, - std::optional max_write = std::nullopt) { + std::optional max_write = std::nullopt, + const char* kind = kv::kind::kSingle) { register_ggml_cache_builders(); kv::CacheGeometry geometry; for (const int window : windows) { @@ -1463,7 +1510,7 @@ InstalledCache install_cache( .initial_capacity = initial_capacity, .max_write = max_write}; auto built = kv::CacheFactory::global().build( - kGgmlCacheBackendId, kv::kind::kSingle, geometry, cfg); + kGgmlCacheBackendId, kind, geometry, cfg); if (!built.ok()) { throw std::runtime_error("failed to build the test cache"); } @@ -1820,6 +1867,399 @@ TEST(GgmlEngineTest, OffGraphRingRewindsWithinTheRowsItHolds) { EXPECT_LT(reference.step(*decode, 8, 1, /*seed=*/4), 1e-4f); } +// What a cell cache should hold: K and V [Hkv, D] per (sequence, position), +// a history per sequence that a clone copies, attended causally within the +// sequence by naive softmax attention once per entry of `windows`, over the +// last `window` positions when it is nonzero, and summed. +class CellReference { + public: + explicit CellReference(std::vector windows = {0}) + : windows_(std::move(windows)) {} + + void clone(int32_t source, int32_t target, int upto) { + history_[target].assign( + history_[source].begin(), history_[source].begin() + upto); + } + void rewind(int32_t seq, int position) { + history_[seq].resize(static_cast(position)); + } + + // Runs one token per (sequence, position) through `executable`, after + // `batch` declares their sequences, and returns the largest difference from + // the reference. + float step( + EngineExecutable& executable, + kv::BatchControl& batch, + const std::vector>& tokens, + int seed = 0) { + const auto n = static_cast(tokens.size()); + const int64_t kv_row = kOffGraphKvHeads * kOffGraphDim; + std::vector q(kOffGraphHeads * n * kOffGraphDim); + std::vector k(kOffGraphKvHeads * n * kOffGraphDim); + std::vector v(k.size()); + std::vector positions; + std::vector seq_ids; + for (int64_t t = 0; t < n; ++t) { + const auto [seq, pos] = tokens[t]; + positions.push_back(pos); + seq_ids.push_back(seq); + const double salt = 1.3 * seq + seed; + std::vector>& history = history_[seq]; + if (static_cast(history.size()) != pos) { + throw std::logic_error("reference: positions must continue"); + } + history.emplace_back(2 * kv_row); + for (int64_t h = 0; h < kOffGraphHeads; ++h) { + for (int64_t d = 0; d < kOffGraphDim; ++d) { + q[(h * n + t) * kOffGraphDim + d] = static_cast( + std::sin(0.17 * pos + 0.9 * h - 0.21 * d + salt)); + } + } + 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 * salt)); + const auto value = static_cast( + std::cos(0.2 * pos - 0.5 * h + 0.31 * d + salt) * (1 + pos % 3)); + k[(h * n + t) * kOffGraphDim + d] = key; + v[(h * n + t) * kOffGraphDim + d] = value; + history.back()[h * kOffGraphDim + d] = key; + history.back()[kv_row + h * kOffGraphDim + d] = value; + } + } + } + if (!batch.declare_step(seq_ids)) { + throw std::runtime_error("reference: the cache refused the step"); + } + 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 t = 0; t < n; ++t) { + const auto [seq, pos] = tokens[t]; + const std::vector>& history = history_[seq]; + for (int64_t h = 0; h < kOffGraphHeads; ++h) { + const int64_t kv_head = h / (kOffGraphHeads / kOffGraphKvHeads); + const float* query = &q[(h * n + t) * kOffGraphDim]; + std::vector want(kOffGraphDim, 0.0); + for (const int window : windows_) { + const int64_t first = + window > 0 ? std::max(0, pos + 1 - window) : 0; + std::vector weights(static_cast(pos + 1), 0.0); + double total = 0.0; + for (int64_t s = first; s <= pos; ++s) { + weights[s] = std::exp( + std::inner_product( + query, + query + kOffGraphDim, + &history[s][kv_head * kOffGraphDim], + 0.0) * + kOffGraphScale); + total += weights[s]; + } + for (int64_t d = 0; d < kOffGraphDim; ++d) { + double sum = 0.0; + for (int64_t s = first; s <= pos; ++s) { + sum += + weights[s] * history[s][kv_row + kv_head * kOffGraphDim + d]; + } + want[d] += sum / total; + } + } + for (int64_t d = 0; d < kOffGraphDim; ++d) { + worst = std::max( + worst, + std::abs( + out[(h * n + t) * kOffGraphDim + d] - + static_cast(want[d]))); + } + } + } + return worst; + } + + private: + std::vector windows_; + std::map>> history_; +}; + +InstalledCache install_cell_cache( + int capacity, + int initial_capacity, + ScalarType kv_dtype = ScalarType::Float, + const std::vector& windows = {0}) { + return install_cache( + capacity, + initial_capacity, + kv_dtype, + static_cast(kOffGraphDim), + windows, + std::nullopt, + kv::kind::kBatchedCell); +} + +// Two sequences prefill separately, then decode together in one forward, +// across graph caching modes and both storage types. Storage starts at one +// row, so it regrows as the cells fill. +TEST(GgmlEngineTest, CellCacheServesInterleavedSequences) { + for (const size_t graphs : {size_t{0}, 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_cell_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(2), cache.options); + auto prefill = prefill_context->compile("forward"); + auto decode = decode_context->compile("forward"); + const int32_t a = *cache.batch().seq_new(); + const int32_t b = *cache.batch().seq_new(); + CellReference reference; + EXPECT_LT( + reference.step(*prefill, cache.batch(), {{a, 0}, {a, 1}, {a, 2}}), + tol); + EXPECT_LT( + reference.step(*prefill, cache.batch(), {{b, 0}, {b, 1}, {b, 2}}), + tol); + for (int64_t pos = 3; pos < 12; ++pos) { + SCOPED_TRACE(pos); + EXPECT_LT( + reference.step(*decode, cache.batch(), {{a, pos}, {b, pos}}), tol); + } + EXPECT_EQ(cache.batch().pos(a), 12); + EXPECT_EQ(cache.batch().pos(b), 12); + } + } +} + +// A clone shares its source's cells below the fork, and each continues on its +// own; a rewind drops the positions it passes. +TEST(GgmlEngineTest, CellCacheClonesAndRewindsSequences) { + InstalledCache cache = install_cell_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(2), cache.options); + auto single_context = + make_context(*host, make_update_and_attend_program(1), cache.options); + auto prefill = prefill_context->compile("forward"); + auto decode = decode_context->compile("forward"); + auto single = single_context->compile("forward"); + const int32_t a = *cache.batch().seq_new(); + CellReference reference; + EXPECT_LT( + reference.step(*prefill, cache.batch(), {{a, 0}, {a, 1}, {a, 2}}), 1e-4f); + for (int64_t pos = 3; pos < 6; ++pos) { + EXPECT_LT(reference.step(*single, cache.batch(), {{a, pos}}), 1e-4f); + } + const auto b = cache.batch().seq_clone(a, 4); + ASSERT_TRUE(b.has_value()); + reference.clone(a, *b, 4); + EXPECT_EQ(cache.batch().pos(*b), 4); + for (int64_t i = 0; i < 3; ++i) { + SCOPED_TRACE(i); + EXPECT_LT( + reference.step(*decode, cache.batch(), {{a, 6 + i}, {*b, 4 + i}}), + 1e-4f); + } + ASSERT_TRUE(cache.batch().rewind(a, 3)); + reference.rewind(a, 3); + EXPECT_LT( + reference.step(*single, cache.batch(), {{a, 3}}, /*seed=*/7), 1e-4f); + EXPECT_LT( + reference.step(*decode, cache.batch(), {{a, 4}, {*b, 7}}, /*seed=*/7), + 1e-4f); + // Closing the clone frees its own cells and leaves the source's. + ASSERT_TRUE(cache.batch().seq_rm(*b)); + EXPECT_LT( + reference.step(*single, cache.batch(), {{a, 5}}, /*seed=*/7), 1e-4f); +} + +// A windowed layer keeps every cell and its mask narrows to the window, within +// a prefill wider than the window too. +TEST(GgmlEngineTest, CellCacheMasksSlidingWindows) { + for (const int window : {2, 5}) { + SCOPED_TRACE(window); + InstalledCache cache = + install_cell_cache(64, 2, ScalarType::Float, {window}); + 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(2), cache.options); + auto prefill = prefill_context->compile("forward"); + auto decode = decode_context->compile("forward"); + const int32_t a = *cache.batch().seq_new(); + const int32_t b = *cache.batch().seq_new(); + CellReference reference({window}); + EXPECT_LT( + reference.step(*prefill, cache.batch(), {{a, 0}, {a, 1}, {a, 2}}), + 1e-4f); + EXPECT_LT( + reference.step(*prefill, cache.batch(), {{b, 0}, {a, 3}, {b, 1}}), + 1e-4f); + for (int64_t i = 0; i < 6; ++i) { + EXPECT_LT( + reference.step(*decode, cache.batch(), {{a, 4 + i}, {b, 2 + i}}), + 1e-4f); + } + } +} + +// The cache places only a declared step whose positions continue each +// sequence. +TEST(GgmlEngineTest, CellCacheRefusesUndeclaredSteps) { + InstalledCache cache = install_cell_cache(64, 4); + auto host = GgmlEngineHost::create(); + auto context = + make_context(*host, make_update_and_attend_program(1), cache.options); + auto executable = context->compile("forward"); + const int32_t a = *cache.batch().seq_new(); + CellReference reference; + const std::vector q(kOffGraphHeads * kOffGraphDim); + const std::vector kv_step(kOffGraphKvHeads * kOffGraphDim); + const auto run_at = [&](int64_t position) { + executable->set_input(0, q.data(), q.size(), ScalarType::Float); + executable->set_input(1, kv_step.data(), kv_step.size(), ScalarType::Float); + executable->set_input(2, kv_step.data(), kv_step.size(), ScalarType::Float); + executable->set_input(3, &position, 1, ScalarType::Long); + executable->execute(); + }; + EXPECT_THROW(run_at(0), std::invalid_argument); + ASSERT_TRUE(cache.batch().declare_step({a})); + EXPECT_THROW(run_at(1), std::invalid_argument); + EXPECT_LT(reference.step(*executable, cache.batch(), {{a, 0}}), 1e-4f); +} + +// One forward serves a sliding and a full layer from the same cells. +TEST(GgmlEngineTest, CellCacheServesSlidingAndFullLayersInOneForward) { + InstalledCache cache = install_cell_cache(64, 2, ScalarType::Float, {3, 0}); + auto host = GgmlEngineHost::create(); + auto prefill_context = + make_context(*host, make_layered_attention_program(3, 2), cache.options); + auto decode_context = + make_context(*host, make_layered_attention_program(2, 2), cache.options); + auto prefill = prefill_context->compile("forward"); + auto decode = decode_context->compile("forward"); + const int32_t a = *cache.batch().seq_new(); + const int32_t b = *cache.batch().seq_new(); + CellReference reference({3, 0}); + EXPECT_LT( + reference.step(*prefill, cache.batch(), {{a, 0}, {a, 1}, {a, 2}}), 1e-4f); + EXPECT_LT( + reference.step(*prefill, cache.batch(), {{b, 0}, {a, 3}, {b, 1}}), 1e-4f); + for (int64_t i = 0; i < 5; ++i) { + EXPECT_LT( + reference.step(*decode, cache.batch(), {{a, 4 + i}, {b, 2 + i}}), + 1e-4f); + } +} + +// A full pool refuses a step until a rewind frees cells. +TEST(GgmlEngineTest, CellCacheRefusesStepsPastAFullPool) { + InstalledCache cache = install_cell_cache(4, 1); + 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(2), cache.options); + auto single_context = + make_context(*host, make_update_and_attend_program(1), cache.options); + auto prefill = prefill_context->compile("forward"); + auto decode = decode_context->compile("forward"); + auto single = single_context->compile("forward"); + const int32_t a = *cache.batch().seq_new(); + const int32_t b = *cache.batch().seq_new(); + CellReference reference; + EXPECT_LT( + reference.step(*prefill, cache.batch(), {{a, 0}, {a, 1}, {a, 2}}), 1e-4f); + EXPECT_THROW( + reference.step(*decode, cache.batch(), {{a, 3}, {b, 0}}), + std::runtime_error); + reference.rewind(a, 3); + reference.rewind(b, 0); + EXPECT_LT(reference.step(*single, cache.batch(), {{b, 0}}), 1e-4f); + EXPECT_EQ(cache.batch().pos(a), 3); + ASSERT_TRUE(cache.batch().rewind(a, 1)); + reference.rewind(a, 1); + EXPECT_LT( + reference.step(*decode, cache.batch(), {{a, 1}, {b, 1}}, /*seed=*/2), + 1e-4f); +} + +// A forward that fails after placing its step, here on a row outside the +// output, gives its cells back, so the sequences continue where they were. +TEST(GgmlEngineTest, FailedCellStepGivesBackItsCells) { + InstalledCache cache = install_cell_cache(64, 4); + auto host = GgmlEngineHost::create(); + auto prefill_context = + make_context(*host, make_update_and_attend_program(3), cache.options); + auto gather_context = make_context( + *host, make_update_and_attend_program(2, 0, true), cache.options); + auto decode_context = + make_context(*host, make_update_and_attend_program(2), cache.options); + auto prefill = prefill_context->compile("forward"); + auto gather = gather_context->compile("forward"); + auto decode = decode_context->compile("forward"); + const int32_t a = *cache.batch().seq_new(); + const int32_t b = *cache.batch().seq_new(); + CellReference reference; + EXPECT_LT( + reference.step(*prefill, cache.batch(), {{a, 0}, {a, 1}, {a, 2}}), 1e-4f); + const std::vector q(kOffGraphHeads * 2 * kOffGraphDim); + const std::vector kv_step(kOffGraphKvHeads * 2 * kOffGraphDim); + const std::vector positions = {3, 0}; + const int64_t row = 2 * kOffGraphHeads; + ASSERT_TRUE(cache.batch().declare_step({a, b})); + 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, positions.data(), positions.size(), ScalarType::Long); + gather->set_input(4, &row, 1, ScalarType::Long); + EXPECT_THROW(gather->execute(), std::out_of_range); + EXPECT_EQ(cache.batch().pos(a), 3); + EXPECT_EQ(cache.batch().pos(b), 0); + EXPECT_LT(reference.step(*decode, cache.batch(), {{a, 3}, {b, 0}}), 1e-4f); +} + +// A forward ends its step, so running it again without a declaration is +// refused rather than re-served, and cannot undo the step that ran. +TEST(GgmlEngineTest, CellStepEndsWithItsForward) { + InstalledCache cache = install_cell_cache(64, 4); + auto host = GgmlEngineHost::create(); + auto gather_context = make_context( + *host, make_update_and_attend_program(1, 0, true), cache.options); + auto gather = gather_context->compile("forward"); + const int32_t a = *cache.batch().seq_new(); + const std::vector q(kOffGraphHeads * kOffGraphDim); + const std::vector kv_step(kOffGraphKvHeads * kOffGraphDim); + const int64_t position = 0; + int64_t row = 0; + 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); + ASSERT_TRUE(cache.batch().declare_step({a})); + gather->execute(); + row = kOffGraphHeads; + gather->set_input(4, &row, 1, ScalarType::Long); + EXPECT_THROW(gather->execute(), std::invalid_argument); + EXPECT_EQ(cache.batch().pos(a), 1); +} + TEST(GgmlEngineTest, OffGraphAttentionNeedsAMatchingCache) { auto host = GgmlEngineHost::create(); // No cache option. diff --git a/extension/llm/cache/cell_cache.cpp b/extension/llm/cache/cell_cache.cpp index 907bdc000ef..5190e2ac224 100644 --- a/extension/llm/cache/cell_cache.cpp +++ b/extension/llm/cache/cell_cache.cpp @@ -208,6 +208,27 @@ CellCache::place_step(int layer, const int32_t* positions, int length) { return &step_for(windows_[layer]); } +void CellCache::end_step() { + invalidate_steps(); +} + +void CellCache::abandon_step() { + if (!placed_) { + return; + } + // Each sequence's first position in the step. + std::map starts; + for (size_t i = 0; i < step_seq_ids_.size(); ++i) { + const auto [it, inserted] = + starts.try_emplace(step_seq_ids_[i], step_pos_[i]); + it->second = std::min(it->second, step_pos_[i]); + } + for (const auto& [seq_id, start] : starts) { + drop_from(seq_id, start); + } + invalidate_steps(); +} + // -- internals --------------------------------------------------------------- uint64_t CellCache::bit(int32_t seq_id) { diff --git a/extension/llm/cache/cell_cache.h b/extension/llm/cache/cell_cache.h index 4ab8365dff9..a47ac17d04a 100644 --- a/extension/llm/cache/cell_cache.h +++ b/extension/llm/cache/cell_cache.h @@ -96,6 +96,16 @@ class ET_EXPERIMENTAL CellCache : public Cache, const CellStep* place_step(int layer, const int32_t* positions, int length) override; + // -- Forward outcome (the byte layer) ----------------------------------- + + // Ends the placed step once its forward has run, so a later forward + // declares again instead of re-serving it. + void end_step(); + // Ends the placed step after its forward failed: each of its sequences + // drops back to its first position in the step, giving back the cells the + // step claimed. Does nothing without a placed step. + void abandon_step(); + protected: void* face(FaceId id) override { return expose(this, id); diff --git a/extension/llm/cache/test/cache_test.cpp b/extension/llm/cache/test/cache_test.cpp index c073428d1ed..8d51c4f9359 100644 --- a/extension/llm/cache/test/cache_test.cpp +++ b/extension/llm/cache/test/cache_test.cpp @@ -684,6 +684,34 @@ TEST_F(CacheTest, CellVerbBetweenLayersInvalidatesTheStep) { EXPECT_EQ(c.place(1, {2}), nullptr); } +TEST_F(CacheTest, CellEndedStepIsNotReServed) { + Cells c(16); + const int32_t s0 = c.seq_new(); + ASSERT_NE(c.step(flatten_step({{s0, 0, 2}})), nullptr); + ASSERT_NE(c.place(1, {0, 1}), nullptr); // another layer of the same forward + + c.cache.end_step(); + EXPECT_EQ(c.place(0, {0, 1}), nullptr); // a later forward must declare + c.cache.abandon_step(); // nothing placed, so nothing to give back + EXPECT_EQ(c.ctl->pos(s0), 2); +} + +TEST_F(CacheTest, CellAbandonedStepGivesBackItsCells) { + Cells c(16); + const int32_t s0 = c.seq_new(); + const int32_t s1 = c.seq_new(); + c.step(flatten_step({{s0, 0, 3}})); + c.cache.end_step(); + ASSERT_NE(c.step(flatten_step({{s0, 3, 1}, {s1, 0, 2}})), nullptr); + + c.cache.abandon_step(); + EXPECT_EQ(c.ctl->pos(s0), 3); + EXPECT_EQ(c.ctl->pos(s1), 0); + EXPECT_EQ(c.cache.free_cells(), 13); + // The same step runs again once declared. + EXPECT_NE(c.step(flatten_step({{s0, 3, 1}, {s1, 0, 2}})), nullptr); +} + TEST_F(CacheTest, CellRewindDropsTheTail) { Cells c(16); const int32_t s0 = c.seq_new();