diff --git a/CMakeLists.txt b/CMakeLists.txt index bfb5600fea7..5d7ac4d7016 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -997,7 +997,9 @@ if(EXECUTORCH_BUILD_EXTENSION_LLM) add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/llm/cache) list(APPEND _executorch_extensions extension_llm_cache) add_subdirectory(${CMAKE_CURRENT_SOURCE_DIR}/extension/llm/batching) - list(APPEND _executorch_extensions extension_llm_batching) + list(APPEND _executorch_extensions extension_llm_batching + extension_llm_batching_cell + ) endif() if(EXECUTORCH_BUILD_EXTENSION_RUNNER_UTIL) diff --git a/extension/llm/batching/CMakeLists.txt b/extension/llm/batching/CMakeLists.txt index 6d00f601b27..0d27e4d20dd 100644 --- a/extension/llm/batching/CMakeLists.txt +++ b/extension/llm/batching/CMakeLists.txt @@ -9,6 +9,10 @@ # scheduler and the executor seam are header-only and free of ExecuTorch runtime # types; the runner owns a thread, so this is a static library rather than an # INTERFACE target. +# +# extension_llm_batching_cell is a separate target because it implements that +# seam against a program and a KV cache, and so carries the runtime types the +# seam itself is kept clear of. if(NOT EXECUTORCH_ROOT) set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../../..) @@ -26,8 +30,22 @@ target_compile_options(extension_llm_batching PUBLIC ${_common_compile_options}) find_package(Threads REQUIRED) target_link_libraries(extension_llm_batching PUBLIC Threads::Threads) +add_library(extension_llm_batching_cell cell_executor.cpp) +target_link_libraries( + extension_llm_batching_cell + PUBLIC extension_llm_batching extension_llm_cache extension_module + extension_tensor + PRIVATE extension_llm_sampler +) +target_include_directories( + extension_llm_batching_cell PUBLIC ${_common_include_directories} +) +target_compile_options( + extension_llm_batching_cell PUBLIC ${_common_compile_options} +) + install( - TARGETS extension_llm_batching + TARGETS extension_llm_batching extension_llm_batching_cell EXPORT ExecuTorchTargets DESTINATION ${CMAKE_INSTALL_LIBDIR} INCLUDES diff --git a/extension/llm/batching/cell_executor.cpp b/extension/llm/batching/cell_executor.cpp new file mode 100644 index 00000000000..d581f132c7d --- /dev/null +++ b/extension/llm/batching/cell_executor.cpp @@ -0,0 +1,402 @@ +/* + * 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 +#include + +namespace executorch { +namespace extension { +namespace llm { +namespace batching { + +using ::executorch::extension::make_tensor_ptr; +using ::executorch::runtime::Error; + +namespace { + +// The backend-load option the delegate resolves the cache through; the +// registry's rendezvous convention, not a per-backend name. +constexpr char kCacheKeyOption[] = "cache_key"; + +std::uint64_t nondeterministic_seed() { + std::random_device device; + return (static_cast(device()) << 32) ^ device(); +} + +} // namespace + +std::optional build_step( + cache::BatchControl& ctl, + const BatchInput& batch, + const std::unordered_map& sessions, + int max_session_tokens) { + Step step; + const std::size_t total = batch.size(); + step.tokens.reserve(total); + step.positions.reserve(total); + step.logit_indices.reserve(batch.inputs.size()); + + std::vector seq_ids; + seq_ids.reserve(total); + // Truncations the batch asks for, held until every input has been checked. + std::vector> rewinds; + // Where each sequence stands mid-batch: the cache still reports what it held + // before the step, so the batch's own writes live here. + std::unordered_map cursor; + + for (const Input& input : batch.inputs) { + const auto seq_it = sessions.find(input.sid); + if (seq_it == sessions.end()) { + ET_LOG(Error, "build_step: session %" PRId64 " is not open", input.sid); + return std::nullopt; + } + const std::int32_t seq = seq_it->second.seq; + if (input.size == 0 || !input.tokens || + input.offset + input.size > input.tokens->size()) { + ET_LOG( + Error, + "build_step: session %" PRId64 " gave a slice its tokens do not hold", + input.sid); + return std::nullopt; + } + + const std::int64_t start = static_cast(input.position) + + static_cast(input.offset); + int& at = cursor.try_emplace(seq, ctl.next_pos(seq)).first->second; + if (start > at) { + // Positions nothing attended, and nothing later reaches back to fill. + ET_LOG( + Error, + "build_step: session %" PRId64 " starts at %" PRId64 + " over a sequence holding %d", + input.sid, + start, + at); + return std::nullopt; + } + if (start < at) { + if (start == 0) { + // Emptying a sequence hands its id back, and the step names it. + ET_LOG( + Error, + "build_step: session %" PRId64 " reopens from the start", + input.sid); + return std::nullopt; + } + rewinds.emplace_back(seq, static_cast(start)); + at = static_cast(start); + } + + const std::int64_t end = start + static_cast(input.size); + if (end > max_session_tokens) { + ET_LOG( + Error, + "build_step: session %" PRId64 " reaches %" PRId64 " of %d cells", + input.sid, + end, + max_session_tokens); + return std::nullopt; + } + + const Token* slice = input.tokens->data() + input.offset; + step.tokens.insert(step.tokens.end(), slice, slice + input.size); + for (std::size_t k = 0; k < input.size; ++k) { + step.positions.push_back(start + static_cast(k)); + } + seq_ids.insert(seq_ids.end(), input.size, seq); + at = static_cast(end); + step.logit_indices.push_back( + input.produce_output ? static_cast(step.tokens.size()) - 1 : -1); + } + + for (const auto& [seq, from] : rewinds) { + if (!ctl.seq_rm(seq, from, std::nullopt)) { + ET_LOG(Error, "build_step: sequence %d would not truncate", seq); + return std::nullopt; + } + } + // After the truncations, so the cells they freed count toward admission. + if (!ctl.declare_step(seq_ids)) { + ET_LOG(Error, "build_step: the cache turned the step down"); + return std::nullopt; + } + return step; +} + +CellExecutor::CellExecutor( + std::unique_ptr module, + std::shared_ptr cache, + std::unique_ptr session, + int max_sessions, + int max_session_tokens, + std::string backend_id, + std::string method, + std::int32_t vocab_size) + : session_(std::move(session)), + cache_(std::move(cache)), + module_(std::move(module)), + ctl_(cache_->as_batch_control()), + max_sessions_(max_sessions), + max_session_tokens_(max_session_tokens), + backend_id_(std::move(backend_id)), + method_(std::move(method)), + vocab_size_(vocab_size) {} + +CellExecutor::~CellExecutor() = default; + +std::unique_ptr CellExecutor::create( + std::unique_ptr module, + int max_sessions, + int max_session_tokens, + int kv_dtype, + int initial_capacity, + std::string method) { + if (module == nullptr) { + ET_LOG(Error, "CellExecutor: no program"); + return nullptr; + } + if (max_sessions <= 0 || max_session_tokens <= 0) { + ET_LOG(Error, "CellExecutor: session limits must be positive"); + return nullptr; + } + if (module->load() != Error::Ok) { // a no-op once the caller has loaded it + ET_LOG(Error, "CellExecutor: the program did not load"); + return nullptr; + } + + auto cfg = cache::et::config_from_program(*module); + if (!cfg.ok()) { + return nullptr; + } + cfg->capacity = max_sessions * max_session_tokens; + cfg->kv_dtype = kv_dtype; + if (initial_capacity >= 0) { + cfg->initial_capacity = initial_capacity; + } + if (!cache::valid(*cfg)) { + ET_LOG(Error, "CellExecutor: the program's layout is unusable"); + return nullptr; + } + + const auto meta = module->method_meta(method); + if (!meta.ok()) { + ET_LOG(Error, "CellExecutor: %s has no metadata", method.c_str()); + return nullptr; + } + + std::string backend_id; + for (std::size_t i = 0; i < meta->num_backends(); ++i) { + const auto name = meta->get_backend_name(i); + if (!name.ok()) { + ET_LOG(Error, "CellExecutor: %s has an unnamed delegate", method.c_str()); + return nullptr; + } + if (backend_id.empty()) { + backend_id = name.get(); + } else if (backend_id != name.get()) { + ET_LOG( + Error, + "CellExecutor: %s spans more than one backend, so which holds the " + "cache is ambiguous", + method.c_str()); + return nullptr; + } + } + if (backend_id.empty()) { + ET_LOG(Error, "CellExecutor: %s delegates to nothing", method.c_str()); + return nullptr; + } + + auto built = + cache::CacheBuilderRegistry::global().build(backend_id, "cell", *cfg); + if (!built.ok()) { + ET_LOG( + Error, + "CellExecutor: no cell cache for backend %s", + backend_id.c_str()); + return nullptr; + } + std::shared_ptr cache = built.get(); + if (cache->as_batch_control() == nullptr) { + ET_LOG(Error, "CellExecutor: the cache carries no sequence identity"); + return nullptr; + } + + if (meta->num_outputs() == 0) { + ET_LOG(Error, "CellExecutor: %s publishes no outputs", method.c_str()); + return nullptr; + } + const auto logits_info = meta->output_tensor_meta(0); + if (!logits_info.ok() || logits_info->sizes().empty()) { + ET_LOG(Error, "CellExecutor: %s has no logits shape", method.c_str()); + return nullptr; + } + const auto logits_sizes = logits_info->sizes(); + + auto session = + std::make_unique(cache::make_unique_key(), cache); + // The delegate resolves the cache from this key while the method loads. + char key[::executorch::runtime::kMaxOptionKeyLength] = {}; + std::memcpy(key, kCacheKeyOption, sizeof(kCacheKeyOption) - 1); + ::executorch::runtime::BackendOptions<1> options; + ::executorch::runtime::LoadBackendOptionsMap options_map; + if (options.set_option(key, session->key().c_str()) != Error::Ok || + options_map.set_options(backend_id.c_str(), options.view()) != + Error::Ok) { + ET_LOG(Error, "CellExecutor: could not name the cache to the backend"); + return nullptr; + } + if (module->load_method( + method, + /*planned_memory=*/nullptr, + /*event_tracer=*/nullptr, + &options_map) != Error::Ok) { + ET_LOG(Error, "CellExecutor: could not load %s", method.c_str()); + return nullptr; + } + + return std::unique_ptr(new CellExecutor( + std::move(module), + std::move(cache), + std::move(session), + max_sessions, + max_session_tokens, + std::move(backend_id), + std::move(method), + logits_sizes[logits_sizes.size() - 1])); +} + +std::optional CellExecutor::open_session() { + if (static_cast(sessions_.size()) >= max_sessions_) { + return std::nullopt; + } + const std::optional seq = ctl_->seq_new(); + if (!seq) { + return std::nullopt; + } + const SessionId session = next_session_++; + sessions_.emplace(session, SessionInfo{*seq, nullptr}); + return session; +} + +void CellExecutor::close_session(SessionId session) { + const auto it = sessions_.find(session); + if (it == sessions_.end()) { + return; + } + // Frees the cells and hands the sequence id back. The session id is not. + ctl_->seq_rm(it->second.seq, 0, std::nullopt); + sessions_.erase(it); +} + +void CellExecutor::set_sampling( + SessionId session, + const SamplingParams& params, + std::optional seed) { + const auto it = sessions_.find(session); + if (it == sessions_.end()) { + return; + } + // One sampler per generation, carrying its own generator state from here on. + it->second.sampler = std::make_unique( + vocab_size_, + params.temperature, + params.top_p, + seed.value_or(nondeterministic_seed())); + it->second.sampler->set_topk(params.top_k); +} + +bool CellExecutor::execute(const BatchInput& batch, BatchOutput& out) { + out.outputs.clear(); + out.outputs.resize(batch.inputs.size()); + + const std::optional step = + build_step(*ctl_, batch, sessions_, max_session_tokens_); + if (!step) { + return false; + } + + auto tokens = + make_tensor_ptr({1, static_cast(step->tokens.size())}, step->tokens); + auto positions = make_tensor_ptr( + {static_cast(step->positions.size())}, step->positions); + auto result = module_->execute(method_, {tokens, positions}); + if (!result.ok()) { + ET_LOG( + Error, + "CellExecutor: %s failed with 0x%x", + method_.c_str(), + static_cast(result.error())); + return false; + } + if (result->empty() || !result->at(0).isTensor()) { + ET_LOG(Error, "CellExecutor: %s returned no logits", method_.c_str()); + return false; + } + // Non-const: the sampler reduces each row in place. Each is read once. + auto logits = result->at(0).toTensor(); + + for (std::size_t i = 0; i < batch.inputs.size(); ++i) { + const int row = step->logit_indices[i]; + if (row < 0) { + continue; // a chunk whose prediction is discarded + } + const SessionId session = batch.inputs[i].sid; + const std::optional token = sample_row(logits, row, session); + if (!token) { + return false; + } + out.outputs[i] = Output{session, {*token}}; + } + return true; +} + +std::optional CellExecutor::sample_row( + ::executorch::aten::Tensor& logits, + int row, + SessionId session) { + const auto it = sessions_.find(session); + if (it == sessions_.end() || it->second.sampler == nullptr) { + ET_LOG( + Error, + "CellExecutor: session %" PRId64 " has no sampling policy", + session); + return std::nullopt; + } + if (row >= logits.numel() / vocab_size_) { + ET_LOG(Error, "CellExecutor: logits hold no row %d", row); + return std::nullopt; + } + // A one-row view over the model's own output: sample_from_logits reduces in + // place and reads the last dimension. + auto one_row = make_tensor_ptr( + {vocab_size_}, + static_cast(logits.mutable_data_ptr()) + + static_cast(row) * vocab_size_ * + ::executorch::runtime::elementSize(logits.scalar_type()), + logits.scalar_type()); + return static_cast(sample_from_logits(*one_row, *it->second.sampler)); +} + +} // namespace batching +} // namespace llm +} // namespace extension +} // namespace executorch diff --git a/extension/llm/batching/cell_executor.h b/extension/llm/batching/cell_executor.h new file mode 100644 index 00000000000..4f2bfbc90f6 --- /dev/null +++ b/extension/llm/batching/cell_executor.h @@ -0,0 +1,149 @@ +/* + * 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 + +// An Executor over the cell KV cache. A session is one cache sequence, a batch +// is one forward carrying every input's tokens end to end on a single axis, +// and the cache's mask is what keeps the sequences apart. +// +// The backend id, cache kind, and option key are configuration, so any backend +// registering a cell cache is served. + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +namespace executorch { +namespace extension { +namespace llm { + +class Sampler; + +namespace batching { + +namespace cache = ::executorch::extension::llm::cache; + +// A session's cache sequence and the sampler its generation draws from. +struct SessionInfo { + std::int32_t seq; + std::unique_ptr sampler; +}; + +// One forward's inputs, flattened across the batch. `tokens` and `positions` +// are filled in one pass so entry i of each names the same token, which is how +// the cache pairs them when it places cells and builds the mask. +struct Step { + std::vector tokens; + std::vector positions; + // One entry per input: the logits row it draws from, or -1 when it produces + // none. An input of any width contributes one, since only its last row + // predicts a token the session does not already hold. + std::vector logit_indices; +}; + +// Flatten the batch, truncate whatever it reopens, and declare it to the cache. +// `ctl` does not move as inputs are laid down, so a per-sequence cursor carries +// the batch's own writes -- consecutive chunks of one prompt abut, and only the +// first can reopen committed ground. +// +// Every input is checked before any is truncated, so a batch refused on its +// contents leaves the cache untouched. +// +// nullopt = an input names a session not in `seqs`, starts past the end of its +// sequence, carries it past `max_session_tokens`, reopens from the start, or +// the cache turned the declaration down. +std::optional build_step( + cache::BatchControl& ctl, + const BatchInput& batch, + const std::unordered_map& sessions, + int max_session_tokens); + +class CellExecutor : public Executor { + public: + ~CellExecutor() override; + + // Builds the cell cache from the layout `module` publishes, binds it to the + // backend, and loads the method. The program must be loaded; its method must + // not be, since the delegate resolves the cache while that load runs. + // + // The table holds `max_sessions` sessions of `max_session_tokens` cells. A + // short table refuses the whole batch and takes every session in it down, so + // capacity is that product exactly and open_session() holds the count -- + // exhaustion is kept unreachable rather than handled. + // + // `kv_dtype` is the ET ScalarType the cache stores K/V in; a negative + // `initial_capacity` leaves the pools to grow from their own default. + // + // The backend is read from the program: the executor drives one cache, so + // the method's attention must be delegated to a single backend. + // + // nullptr = unusable limits, a program publishing no KV layout, a method + // spanning several backends, or no cell cache registered for the one it + // names. + static std::unique_ptr create( + std::unique_ptr module, + int max_sessions, + int max_session_tokens, + int kv_dtype, + int initial_capacity = -1, + std::string method = "forward"); + + std::optional open_session() override; + void close_session(SessionId session) override; + void set_sampling( + SessionId session, + const SamplingParams& params, + std::optional seed) override; + bool execute(const BatchInput& batch, BatchOutput& out) override; + + private: + CellExecutor( + std::unique_ptr module, + std::shared_ptr cache, + std::unique_ptr session, + int max_sessions, + int max_session_tokens, + std::string backend_id, + std::string method, + std::int32_t vocab_size); + + // Draw the token an input produced from its row of `logits`, which the + // session's sampler consumes in place. + std::optional + sample_row(::executorch::aten::Tensor& logits, int row, SessionId session); + + // Ordered so the module dies first, releasing the delegate that resolved the + // cache before the registry entry naming it goes. + std::unique_ptr session_; + std::shared_ptr cache_; + std::unique_ptr module_; + cache::BatchControl* ctl_; + int max_sessions_; + int max_session_tokens_; + std::string backend_id_; + std::string method_; + // The method's logits width, so a sampler can be built by its policy. + std::int32_t vocab_size_; + + SessionId next_session_ = 1; // never reused, unlike the cache's sequence ids + std::unordered_map sessions_; +}; + +} // namespace batching +} // namespace llm +} // namespace extension +} // namespace executorch diff --git a/extension/llm/cache/CMakeLists.txt b/extension/llm/cache/CMakeLists.txt index 8e3be66f5d7..f7622a428a4 100644 --- a/extension/llm/cache/CMakeLists.txt +++ b/extension/llm/cache/CMakeLists.txt @@ -15,7 +15,9 @@ if(NOT EXECUTORCH_ROOT) endif() add_library(extension_llm_cache cache_registry.cpp cell_cache.cpp) -target_link_libraries(extension_llm_cache executorch_core) +# cache_et.h reads a program's published KV layout, so the ET adapter needs +# Module even though the neutral core (cache.h) stays ET-independent. +target_link_libraries(extension_llm_cache executorch_core extension_module) target_include_directories( extension_llm_cache PUBLIC ${_common_include_directories} ) diff --git a/extension/llm/cache/cache_et.h b/extension/llm/cache/cache_et.h index 1c157f8ed60..24a93efd64a 100644 --- a/extension/llm/cache/cache_et.h +++ b/extension/llm/cache/cache_et.h @@ -17,8 +17,10 @@ // adapter.) #include +#include #include +#include #include #include @@ -57,6 +59,68 @@ inline Error rewind(SequenceControl& control, int new_len) { return Error::Ok; } +// A CacheConfig's geometry is a property of the model, not a choice: how many +// caches, and each one's heads, head dim, and attention window. An off-graph +// export publishes those as constant methods, which carry no delegate, so +// reading them needs only the program loaded. +// +// The sizing left unset -- capacity, kv_dtype, initial_capacity -- is the +// caller's policy, and the same program runs under any of it. +// +// InvalidArgument if the program publishes no layout, so it is not an +// off-graph model. +inline Result config_from_program(Module& module) { + const auto read_int = [&module](const char* name) -> std::optional { + const auto r = module.execute(name); + if (!r.ok() || r->empty() || !r->at(0).isInt()) { + return std::nullopt; + } + return r->at(0).toInt(); + }; + const auto read_ints = + [&module](const char* name) -> std::optional> { + const auto r = module.execute(name); + if (!r.ok() || r->empty() || !r->at(0).isTensor()) { + return std::nullopt; + } + const auto t = r->at(0).toTensor(); + if (t.scalar_type() != ::executorch::aten::ScalarType::Int) { + return std::nullopt; + } + const int32_t* p = t.const_data_ptr(); + return std::vector(p, p + t.numel()); + }; + + const auto n_caches = read_int("get_n_caches"); + const auto kv_heads = read_ints("get_kv_heads"); + const auto head_dims = read_ints("get_head_dims"); + const auto windows = read_ints("get_windows"); + ET_CHECK_OR_RETURN_ERROR( + n_caches && kv_heads && head_dims && windows, + InvalidArgument, + "cache: the program publishes no KV layout"); + const auto n = static_cast(*n_caches); + ET_CHECK_OR_RETURN_ERROR( + kv_heads->size() == n && head_dims->size() == n && windows->size() == n, + InvalidArgument, + "cache: the published KV layout names %zu caches inconsistently", + n); + + CacheConfig cfg{}; + cfg.n_layers = static_cast(n); + cfg.layers.reserve(n); + for (size_t l = 0; l < n; ++l) { + LayerConfig lc{}; + lc.n_kv_heads = (*kv_heads)[l]; + lc.head_dim = (*head_dims)[l]; + lc.policy = (*windows)[l] > 0 + ? LayerPolicy{LayerPolicy::Kind::Ring, (*windows)[l]} + : LayerPolicy{LayerPolicy::Kind::Flat, 0}; + cfg.layers.push_back(lc); + } + return cfg; +} + } // namespace et } // namespace cache } // namespace llm