From 3e31cf1b8a5b701f9a37127aa7457614bb088535 Mon Sep 17 00:00:00 2001 From: kiymetakdemir Date: Thu, 27 Aug 2026 12:21:45 -0700 Subject: [PATCH 1/3] Add a cell-cache executor for batched generation --- CMakeLists.txt | 4 +- extension/llm/batching/CMakeLists.txt | 20 +- extension/llm/batching/cell_executor.cpp | 406 +++++++++++++++++++++++ extension/llm/batching/cell_executor.h | 148 +++++++++ 4 files changed, 576 insertions(+), 2 deletions(-) create mode 100644 extension/llm/batching/cell_executor.cpp create mode 100644 extension/llm/batching/cell_executor.h 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..5f1118d2478 --- /dev/null +++ b/extension/llm/batching/cell_executor.cpp @@ -0,0 +1,406 @@ +/* + * 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 + +namespace executorch { +namespace extension { +namespace llm { +namespace batching { + +using ::executorch::extension::make_tensor_ptr; +using ::executorch::runtime::Error; + +namespace { + +// Spread one seed over positions so a session's neighbouring tokens draw +// unrelated streams. +std::uint64_t mix(std::uint64_t seed, std::int64_t position) { + std::uint64_t x = + seed + 0x9e3779b97f4a7c15ULL * static_cast(position + 1); + x = (x ^ (x >> 30)) * 0xbf58476d1ce4e5b9ULL; + x = (x ^ (x >> 27)) * 0x94d049bb133111ebULL; + return x ^ (x >> 31); +} + +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& seqs, + 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 = seqs.find(input.sid); + if (seq_it == seqs.end()) { + ET_LOG(Error, "build_step: session %" PRId64 " is not open", input.sid); + return std::nullopt; + } + const std::int32_t seq = seq_it->second; + 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( + Config config, + std::unique_ptr module, + std::shared_ptr cache, + std::unique_ptr session) + : config_(std::move(config)), + session_(std::move(session)), + cache_(std::move(cache)), + module_(std::move(module)), + ctl_(cache_->as_batch_control()) {} + +CellExecutor::~CellExecutor() = default; + +std::unique_ptr CellExecutor::create( + std::unique_ptr module, + Config config) { + if (module == nullptr) { + ET_LOG(Error, "CellExecutor: no program"); + return nullptr; + } + if (config.max_sessions <= 0 || config.max_session_tokens <= 0) { + ET_LOG(Error, "CellExecutor: session limits must be positive"); + return nullptr; + } + // A batch cannot be refused in part, so the table must hold every session at + // its bound rather than discover it is short mid-run. + const std::int64_t needed = static_cast(config.max_sessions) * + config.max_session_tokens; + if (needed > config.cache.capacity) { + ET_LOG( + Error, + "CellExecutor: %d sessions of %d cells need %" PRId64 + ", capacity is %d", + config.max_sessions, + config.max_session_tokens, + needed, + config.cache.capacity); + return nullptr; + } + if (!cache::valid(config.cache)) { + ET_LOG(Error, "CellExecutor: invalid cache config"); + 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 built = cache::CacheBuilderRegistry::global().build( + config.backend_id, config.cache_kind, config.cache); + if (!built.ok()) { + ET_LOG( + Error, + "CellExecutor: no %s cache for backend %s", + config.cache_kind.c_str(), + config.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; + } + + // Checked here rather than at the deferred load: it needs only the config. + if (config.cache_key_option.size() >= + ::executorch::runtime::kMaxOptionKeyLength) { + ET_LOG( + Error, + "CellExecutor: option key %s is too long", + config.cache_key_option.c_str()); + return nullptr; + } + + auto session = + std::make_unique(cache::make_unique_key(), cache); + return std::unique_ptr(new CellExecutor( + std::move(config), + std::move(module), + std::move(cache), + std::move(session))); +} + +bool CellExecutor::ensure_method_loaded() { + if (method_loaded_) { + return true; + } + // The key is taken as a fixed-width array; create() bounded its length. + char key[::executorch::runtime::kMaxOptionKeyLength] = {}; + std::memcpy( + key, config_.cache_key_option.data(), config_.cache_key_option.size()); + + ::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(config_.backend_id.c_str(), options.view()) != + Error::Ok) { + ET_LOG(Error, "CellExecutor: could not name the cache to the backend"); + return false; + } + if (module_->load_method( + config_.method, + /*planned_memory=*/nullptr, + /*event_tracer=*/nullptr, + &options_map) != Error::Ok) { + ET_LOG(Error, "CellExecutor: could not load %s", config_.method.c_str()); + return false; + } + method_loaded_ = true; + return true; +} + +std::optional CellExecutor::open_session() { + if (static_cast(seqs_.size()) >= config_.max_sessions) { + return std::nullopt; + } + const std::optional seq = ctl_->seq_new(); + if (!seq) { + return std::nullopt; + } + const SessionId session = next_session_++; + seqs_.emplace(session, *seq); + return session; +} + +void CellExecutor::close_session(SessionId session) { + const auto it = seqs_.find(session); + if (it == seqs_.end()) { + return; + } + // Frees the cells and hands the sequence id back. The session id is not. + ctl_->seq_rm(it->second, 0, std::nullopt); + seqs_.erase(it); + sampling_.erase(session); +} + +void CellExecutor::set_sampling( + SessionId session, + const SamplingParams& params, + std::optional seed) { + if (seqs_.count(session) == 0) { + return; + } + sampling_.insert_or_assign( + session, Sampling{params, seed.value_or(nondeterministic_seed())}); +} + +bool CellExecutor::execute(const BatchInput& batch, BatchOutput& out) { + out.outputs.clear(); + out.outputs.resize(batch.inputs.size()); + + // Ahead of the step: a refused load must claim no cell. + if (!ensure_method_loaded()) { + return false; + } + + const std::optional step = + build_step(*ctl_, batch, seqs_, config_.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); + const auto result = module_->execute(config_.method, {tokens, positions}); + if (!result.ok()) { + ET_LOG( + Error, + "CellExecutor: %s failed with 0x%x", + config_.method.c_str(), + static_cast(result.error())); + return false; + } + if (result->empty() || !result->at(0).isTensor()) { + ET_LOG( + Error, "CellExecutor: %s returned no logits", config_.method.c_str()); + return false; + } + const 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; + // The drawn token lands one past the row that predicted it. + const std::optional token = + sample_row(logits, row, session, step->positions[row] + 1); + if (!token) { + return false; + } + out.outputs[i] = Output{session, {*token}}; + } + return true; +} + +std::optional CellExecutor::sample_row( + const ::executorch::aten::Tensor& logits, + int row, + SessionId session, + std::int64_t position) const { + const auto policy = sampling_.find(session); + if (policy == sampling_.end()) { + ET_LOG( + Error, + "CellExecutor: session %" PRId64 " has no sampling policy", + session); + return std::nullopt; + } + const SamplingParams& params = policy->second.params; + const auto vocab = logits.size(logits.dim() - 1); + if (row >= logits.numel() / vocab) { + ET_LOG(Error, "CellExecutor: logits hold no row %d", row); + return std::nullopt; + } + + std::optional drawn; + struct { + [[noreturn]] void fail(Error) { + ET_CHECK_MSG(false, "CellExecutor: unsupported logits dtype"); + } + } ctx; + ET_SWITCH_THREE_TYPES( + Float, + Half, + BFloat16, + logits.scalar_type(), + ctx, + "sample_row", + CTYPE, + [&] { + const CTYPE* begin = logits.const_data_ptr() + row * vocab; + if (params.temperature <= 0.0f) { + drawn = static_cast( + std::max_element(begin, begin + vocab) - begin); + return; + } + // The sampler reduces in place, so copy rather than overwrite the + // model's output. + std::vector scores(begin, begin + vocab); + Sampler sampler( + static_cast(vocab), + params.temperature, + params.top_p, + mix(policy->second.seed, position)); + sampler.set_topk(params.top_k); + drawn = static_cast(sampler.sample(scores.data())); + }); + return drawn; +} + +} // 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..bef3131ebc6 --- /dev/null +++ b/extension/llm/batching/cell_executor.h @@ -0,0 +1,148 @@ +/* + * 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 { +namespace batching { + +namespace cache = ::executorch::extension::llm::cache; + +// 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& seqs, + int max_session_tokens); + +class CellExecutor : public Executor { + public: + struct Config { + cache::CacheConfig cache; + // Sessions open at once, and the cells each may hold. A short table refuses + // the whole batch and takes every session in it down, so create() rejects + // limits the table cannot honor and open_session() holds the count -- + // exhaustion is kept unreachable rather than handled. + int max_sessions = 0; + int max_session_tokens = 0; + std::string backend_id; + std::string cache_kind = "cell"; + // Backend-load option through which the delegate finds the cache. + std::string cache_key_option = "cache_key"; + std::string method = "forward"; + }; + + ~CellExecutor() override; + + // Takes the program loaded but its method not. The cache is sized from a + // layout the program publishes, so the program has to be readable first. + // + // The method is left for the first execute(). A delegate may bind per-thread + // state as it initializes, which happens during that load, and construction + // is the only entry point that does not run on the engine thread. So a + // method that will not load fails a batch rather than this call. + // + // nullptr = limits the cell table cannot honor, or an unregistered + // (backend_id, cache_kind). + static std::unique_ptr create( + std::unique_ptr module, + Config config); + + 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: + struct Sampling { + SamplingParams params; + std::uint64_t seed; + }; + + CellExecutor( + Config config, + std::unique_ptr module, + std::shared_ptr cache, + std::unique_ptr session); + + // Load the method, naming the cache to the backend, on first use. Runs on + // the execute() thread so what the delegate binds there is reachable later. + bool ensure_method_loaded(); + + // Draw from one logits row. Randomness comes from the session's seed and the + // position the token will occupy, so it does not follow how batches formed. + std::optional sample_row( + const ::executorch::aten::Tensor& logits, + int row, + SessionId session, + std::int64_t position) const; + + Config config_; + // 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_; + + bool method_loaded_ = false; + SessionId next_session_ = 1; // never reused, unlike the cache's sequence ids + std::unordered_map seqs_; + std::unordered_map sampling_; +}; + +} // namespace batching +} // namespace llm +} // namespace extension +} // namespace executorch From ef6400a5c2d171635ee7d883ebf7faadd863564a Mon Sep 17 00:00:00 2001 From: kiymetakdemir Date: Fri, 28 Aug 2026 15:34:33 -0700 Subject: [PATCH 2/3] Address review: build the cache from the program, add a start hook --- extension/llm/batching/cell_executor.cpp | 224 +++++++++++------------ extension/llm/batching/cell_executor.h | 99 +++++----- extension/llm/batching/executor.h | 8 + extension/llm/batching/runner.cpp | 5 + extension/llm/cache/CMakeLists.txt | 4 +- extension/llm/cache/cache_et.h | 64 +++++++ 6 files changed, 235 insertions(+), 169 deletions(-) diff --git a/extension/llm/batching/cell_executor.cpp b/extension/llm/batching/cell_executor.cpp index 5f1118d2478..6bb6a614e15 100644 --- a/extension/llm/batching/cell_executor.cpp +++ b/extension/llm/batching/cell_executor.cpp @@ -14,7 +14,9 @@ #include #include +#include #include +#include #include #include #include @@ -30,15 +32,9 @@ using ::executorch::runtime::Error; namespace { -// Spread one seed over positions so a session's neighbouring tokens draw -// unrelated streams. -std::uint64_t mix(std::uint64_t seed, std::int64_t position) { - std::uint64_t x = - seed + 0x9e3779b97f4a7c15ULL * static_cast(position + 1); - x = (x ^ (x >> 30)) * 0xbf58476d1ce4e5b9ULL; - x = (x ^ (x >> 27)) * 0x94d049bb133111ebULL; - return x ^ (x >> 31); -} +// 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; @@ -50,7 +46,7 @@ std::uint64_t nondeterministic_seed() { std::optional build_step( cache::BatchControl& ctl, const BatchInput& batch, - const std::unordered_map& seqs, + const std::unordered_map& sessions, int max_session_tokens) { Step step; const std::size_t total = batch.size(); @@ -67,12 +63,12 @@ std::optional build_step( std::unordered_map cursor; for (const Input& input : batch.inputs) { - const auto seq_it = seqs.find(input.sid); - if (seq_it == seqs.end()) { + 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; + const std::int32_t seq = seq_it->second.seq; if (input.size == 0 || !input.tokens || input.offset + input.size > input.tokens->size()) { ET_LOG( @@ -146,62 +142,68 @@ std::optional build_step( } CellExecutor::CellExecutor( - Config config, std::unique_ptr module, std::shared_ptr cache, - std::unique_ptr session) - : config_(std::move(config)), - session_(std::move(session)), + 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()) {} + 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, - Config config) { + int max_sessions, + int max_session_tokens, + int kv_dtype, + std::string backend_id, + int initial_capacity, + std::string method) { if (module == nullptr) { ET_LOG(Error, "CellExecutor: no program"); return nullptr; } - if (config.max_sessions <= 0 || config.max_session_tokens <= 0) { + if (max_sessions <= 0 || max_session_tokens <= 0) { ET_LOG(Error, "CellExecutor: session limits must be positive"); return nullptr; } - // A batch cannot be refused in part, so the table must hold every session at - // its bound rather than discover it is short mid-run. - const std::int64_t needed = static_cast(config.max_sessions) * - config.max_session_tokens; - if (needed > config.cache.capacity) { - ET_LOG( - Error, - "CellExecutor: %d sessions of %d cells need %" PRId64 - ", capacity is %d", - config.max_sessions, - config.max_session_tokens, - needed, - config.cache.capacity); + 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; } - if (!cache::valid(config.cache)) { - ET_LOG(Error, "CellExecutor: invalid cache config"); + + auto cfg = cache::et::config_from_program(*module); + if (!cfg.ok()) { 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"); + 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; } - auto built = cache::CacheBuilderRegistry::global().build( - config.backend_id, config.cache_kind, config.cache); + auto built = + cache::CacheBuilderRegistry::global().build(backend_id, "cell", *cfg); if (!built.ok()) { ET_LOG( Error, - "CellExecutor: no %s cache for backend %s", - config.cache_kind.c_str(), - config.backend_id.c_str()); + "CellExecutor: no cell cache for backend %s", + backend_id.c_str()); return nullptr; } std::shared_ptr cache = built.get(); @@ -210,48 +212,53 @@ std::unique_ptr CellExecutor::create( return nullptr; } - // Checked here rather than at the deferred load: it needs only the config. - if (config.cache_key_option.size() >= - ::executorch::runtime::kMaxOptionKeyLength) { - ET_LOG( - Error, - "CellExecutor: option key %s is too long", - config.cache_key_option.c_str()); + const auto meta = module->method_meta(method); + if (!meta.ok() || 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); return std::unique_ptr(new CellExecutor( - std::move(config), std::move(module), std::move(cache), - std::move(session))); + std::move(session), + max_sessions, + max_session_tokens, + std::move(backend_id), + std::move(method), + logits_sizes[logits_sizes.size() - 1])); } -bool CellExecutor::ensure_method_loaded() { +bool CellExecutor::start() { if (method_loaded_) { return true; } // The key is taken as a fixed-width array; create() bounded its length. char key[::executorch::runtime::kMaxOptionKeyLength] = {}; - std::memcpy( - key, config_.cache_key_option.data(), config_.cache_key_option.size()); + 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(config_.backend_id.c_str(), options.view()) != + 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 false; } if (module_->load_method( - config_.method, + method_, /*planned_memory=*/nullptr, /*event_tracer=*/nullptr, &options_map) != Error::Ok) { - ET_LOG(Error, "CellExecutor: could not load %s", config_.method.c_str()); + ET_LOG(Error, "CellExecutor: could not load %s", method_.c_str()); return false; } method_loaded_ = true; @@ -259,7 +266,7 @@ bool CellExecutor::ensure_method_loaded() { } std::optional CellExecutor::open_session() { - if (static_cast(seqs_.size()) >= config_.max_sessions) { + if (static_cast(sessions_.size()) >= max_sessions_) { return std::nullopt; } const std::optional seq = ctl_->seq_new(); @@ -267,43 +274,48 @@ std::optional CellExecutor::open_session() { return std::nullopt; } const SessionId session = next_session_++; - seqs_.emplace(session, *seq); + sessions_.emplace(session, SessionInfo{*seq, nullptr}); return session; } void CellExecutor::close_session(SessionId session) { - const auto it = seqs_.find(session); - if (it == seqs_.end()) { + 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, 0, std::nullopt); - seqs_.erase(it); - sampling_.erase(session); + ctl_->seq_rm(it->second.seq, 0, std::nullopt); + sessions_.erase(it); } void CellExecutor::set_sampling( SessionId session, const SamplingParams& params, std::optional seed) { - if (seqs_.count(session) == 0) { + const auto it = sessions_.find(session); + if (it == sessions_.end()) { return; } - sampling_.insert_or_assign( - session, Sampling{params, seed.value_or(nondeterministic_seed())}); + // 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()); - // Ahead of the step: a refused load must claim no cell. - if (!ensure_method_loaded()) { + if (!method_loaded_) { + ET_LOG(Error, "CellExecutor: execute() before start()"); return false; } const std::optional step = - build_step(*ctl_, batch, seqs_, config_.max_session_tokens); + build_step(*ctl_, batch, sessions_, max_session_tokens_); if (!step) { return false; } @@ -312,21 +324,21 @@ bool CellExecutor::execute(const BatchInput& batch, BatchOutput& out) { make_tensor_ptr({1, static_cast(step->tokens.size())}, step->tokens); auto positions = make_tensor_ptr( {static_cast(step->positions.size())}, step->positions); - const auto result = module_->execute(config_.method, {tokens, positions}); + auto result = module_->execute(method_, {tokens, positions}); if (!result.ok()) { ET_LOG( Error, "CellExecutor: %s failed with 0x%x", - config_.method.c_str(), + method_.c_str(), static_cast(result.error())); return false; } if (result->empty() || !result->at(0).isTensor()) { - ET_LOG( - Error, "CellExecutor: %s returned no logits", config_.method.c_str()); + ET_LOG(Error, "CellExecutor: %s returned no logits", method_.c_str()); return false; } - const auto logits = result->at(0).toTensor(); + // 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]; @@ -334,9 +346,7 @@ bool CellExecutor::execute(const BatchInput& batch, BatchOutput& out) { continue; // a chunk whose prediction is discarded } const SessionId session = batch.inputs[i].sid; - // The drawn token lands one past the row that predicted it. - const std::optional token = - sample_row(logits, row, session, step->positions[row] + 1); + const std::optional token = sample_row(logits, row, session); if (!token) { return false; } @@ -346,58 +356,30 @@ bool CellExecutor::execute(const BatchInput& batch, BatchOutput& out) { } std::optional CellExecutor::sample_row( - const ::executorch::aten::Tensor& logits, + ::executorch::aten::Tensor& logits, int row, - SessionId session, - std::int64_t position) const { - const auto policy = sampling_.find(session); - if (policy == sampling_.end()) { + 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; } - const SamplingParams& params = policy->second.params; - const auto vocab = logits.size(logits.dim() - 1); - if (row >= logits.numel() / vocab) { + if (row >= logits.numel() / vocab_size_) { ET_LOG(Error, "CellExecutor: logits hold no row %d", row); return std::nullopt; } - - std::optional drawn; - struct { - [[noreturn]] void fail(Error) { - ET_CHECK_MSG(false, "CellExecutor: unsupported logits dtype"); - } - } ctx; - ET_SWITCH_THREE_TYPES( - Float, - Half, - BFloat16, - logits.scalar_type(), - ctx, - "sample_row", - CTYPE, - [&] { - const CTYPE* begin = logits.const_data_ptr() + row * vocab; - if (params.temperature <= 0.0f) { - drawn = static_cast( - std::max_element(begin, begin + vocab) - begin); - return; - } - // The sampler reduces in place, so copy rather than overwrite the - // model's output. - std::vector scores(begin, begin + vocab); - Sampler sampler( - static_cast(vocab), - params.temperature, - params.top_p, - mix(policy->second.seed, position)); - sampler.set_topk(params.top_k); - drawn = static_cast(sampler.sample(scores.data())); - }); - return drawn; + // 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 diff --git a/extension/llm/batching/cell_executor.h b/extension/llm/batching/cell_executor.h index bef3131ebc6..58f8da83bc3 100644 --- a/extension/llm/batching/cell_executor.h +++ b/extension/llm/batching/cell_executor.h @@ -30,10 +30,19 @@ 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. @@ -60,41 +69,41 @@ struct Step { std::optional build_step( cache::BatchControl& ctl, const BatchInput& batch, - const std::unordered_map& seqs, + const std::unordered_map& sessions, int max_session_tokens); class CellExecutor : public Executor { public: - struct Config { - cache::CacheConfig cache; - // Sessions open at once, and the cells each may hold. A short table refuses - // the whole batch and takes every session in it down, so create() rejects - // limits the table cannot honor and open_session() holds the count -- - // exhaustion is kept unreachable rather than handled. - int max_sessions = 0; - int max_session_tokens = 0; - std::string backend_id; - std::string cache_kind = "cell"; - // Backend-load option through which the delegate finds the cache. - std::string cache_key_option = "cache_key"; - std::string method = "forward"; - }; - ~CellExecutor() override; - // Takes the program loaded but its method not. The cache is sized from a - // layout the program publishes, so the program has to be readable first. + // Builds the cell cache from the layout `module` publishes and binds it to + // the backend. The program must be loaded; its method must not be, since the + // method load is left until the first batch -- a delegate may bind per-thread + // state as it initializes, and construction is the only entry point that does + // not run on the engine thread. + // + // 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. // - // The method is left for the first execute(). A delegate may bind per-thread - // state as it initializes, which happens during that load, and construction - // is the only entry point that does not run on the engine thread. So a - // method that will not load fails a batch rather than this call. + // `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. // - // nullptr = limits the cell table cannot honor, or an unregistered - // (backend_id, cache_kind). + // nullptr = unusable limits, a program publishing no KV layout, or no cell + // cache registered for `backend_id`. static std::unique_ptr create( std::unique_ptr module, - Config config); + int max_sessions, + int max_session_tokens, + int kv_dtype, + std::string backend_id, + int initial_capacity = -1, + std::string method = "forward"); + + // Loads the method, naming the cache to the backend. The delegate binds + // per-thread state as it initializes, which happens during that load. + bool start() override; std::optional open_session() override; void close_session(SessionId session) override; @@ -105,41 +114,37 @@ class CellExecutor : public Executor { bool execute(const BatchInput& batch, BatchOutput& out) override; private: - struct Sampling { - SamplingParams params; - std::uint64_t seed; - }; - CellExecutor( - Config config, std::unique_ptr module, std::shared_ptr cache, - std::unique_ptr session); - - // Load the method, naming the cache to the backend, on first use. Runs on - // the execute() thread so what the delegate binds there is reachable later. - bool ensure_method_loaded(); - - // Draw from one logits row. Randomness comes from the session's seed and the - // position the token will occupy, so it does not follow how batches formed. - std::optional sample_row( - const ::executorch::aten::Tensor& logits, - int row, - SessionId session, - std::int64_t position) const; + 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); - Config config_; // 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_; bool method_loaded_ = false; SessionId next_session_ = 1; // never reused, unlike the cache's sequence ids - std::unordered_map seqs_; - std::unordered_map sampling_; + std::unordered_map sessions_; }; } // namespace batching diff --git a/extension/llm/batching/executor.h b/extension/llm/batching/executor.h index 87a95946889..1b526542664 100644 --- a/extension/llm/batching/executor.h +++ b/extension/llm/batching/executor.h @@ -67,6 +67,14 @@ class Executor { public: virtual ~Executor() = default; + // Bind whatever must live on the thread that will run the batches. Called + // once before any other call, on that thread. + // + // false = the executor cannot run. + virtual bool start() { + return true; + } + // A session to route tasks to. nullopt = at capacity. Every successful id // must be unique for the lifetime of the consuming Runner, even after close. virtual std::optional open_session() = 0; diff --git a/extension/llm/batching/runner.cpp b/extension/llm/batching/runner.cpp index 7a21c3910bd..807337ac7d6 100644 --- a/extension/llm/batching/runner.cpp +++ b/extension/llm/batching/runner.cpp @@ -454,6 +454,11 @@ void RunnerImpl::notify_engine_() { // --- engine thread --------------------------------------------------------- void RunnerImpl::run_() { + // Start before anything is admitted, and on this thread. + if (!executor_.start()) { + lifecycle_.store(Lifecycle::Stopping, std::memory_order_release); + } + while (is_running_()) { process_pending_commands_(); reap_cancelled_(); 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 From cb150089e43a832b263693821d838796eacf1bfe Mon Sep 17 00:00:00 2001 From: kiymetakdemir Date: Fri, 28 Aug 2026 19:00:17 -0700 Subject: [PATCH 3/3] Address review: build the cache from the program, read the backend from it --- extension/llm/batching/cell_executor.cpp | 82 ++++++++++++++---------- extension/llm/batching/cell_executor.h | 22 +++---- extension/llm/batching/executor.h | 8 --- extension/llm/batching/runner.cpp | 5 -- 4 files changed, 57 insertions(+), 60 deletions(-) diff --git a/extension/llm/batching/cell_executor.cpp b/extension/llm/batching/cell_executor.cpp index 6bb6a614e15..d581f132c7d 100644 --- a/extension/llm/batching/cell_executor.cpp +++ b/extension/llm/batching/cell_executor.cpp @@ -167,7 +167,6 @@ std::unique_ptr CellExecutor::create( int max_sessions, int max_session_tokens, int kv_dtype, - std::string backend_id, int initial_capacity, std::string method) { if (module == nullptr) { @@ -197,6 +196,35 @@ std::unique_ptr CellExecutor::create( 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()) { @@ -212,8 +240,7 @@ std::unique_ptr CellExecutor::create( return nullptr; } - const auto meta = module->method_meta(method); - if (!meta.ok() || meta->num_outputs() == 0) { + if (meta->num_outputs() == 0) { ET_LOG(Error, "CellExecutor: %s publishes no outputs", method.c_str()); return nullptr; } @@ -226,43 +253,35 @@ std::unique_ptr CellExecutor::create( auto session = std::make_unique(cache::make_unique_key(), cache); - 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])); -} - -bool CellExecutor::start() { - if (method_loaded_) { - return true; - } - // The key is taken as a fixed-width array; create() bounded its length. + // 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()) != + 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 false; + return nullptr; } - if (module_->load_method( - method_, + 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 false; + ET_LOG(Error, "CellExecutor: could not load %s", method.c_str()); + return nullptr; } - method_loaded_ = true; - return true; + + 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() { @@ -309,11 +328,6 @@ bool CellExecutor::execute(const BatchInput& batch, BatchOutput& out) { out.outputs.clear(); out.outputs.resize(batch.inputs.size()); - if (!method_loaded_) { - ET_LOG(Error, "CellExecutor: execute() before start()"); - return false; - } - const std::optional step = build_step(*ctl_, batch, sessions_, max_session_tokens_); if (!step) { diff --git a/extension/llm/batching/cell_executor.h b/extension/llm/batching/cell_executor.h index 58f8da83bc3..4f2bfbc90f6 100644 --- a/extension/llm/batching/cell_executor.h +++ b/extension/llm/batching/cell_executor.h @@ -76,11 +76,9 @@ class CellExecutor : public Executor { public: ~CellExecutor() override; - // Builds the cell cache from the layout `module` publishes and binds it to - // the backend. The program must be loaded; its method must not be, since the - // method load is left until the first batch -- a delegate may bind per-thread - // state as it initializes, and construction is the only entry point that does - // not run on the engine thread. + // 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 @@ -90,21 +88,20 @@ class CellExecutor : public Executor { // `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. // - // nullptr = unusable limits, a program publishing no KV layout, or no cell - // cache registered for `backend_id`. + // 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, - std::string backend_id, int initial_capacity = -1, std::string method = "forward"); - // Loads the method, naming the cache to the backend. The delegate binds - // per-thread state as it initializes, which happens during that load. - bool start() override; - std::optional open_session() override; void close_session(SessionId session) override; void set_sampling( @@ -142,7 +139,6 @@ class CellExecutor : public Executor { // The method's logits width, so a sampler can be built by its policy. std::int32_t vocab_size_; - bool method_loaded_ = false; SessionId next_session_ = 1; // never reused, unlike the cache's sequence ids std::unordered_map sessions_; }; diff --git a/extension/llm/batching/executor.h b/extension/llm/batching/executor.h index 1b526542664..87a95946889 100644 --- a/extension/llm/batching/executor.h +++ b/extension/llm/batching/executor.h @@ -67,14 +67,6 @@ class Executor { public: virtual ~Executor() = default; - // Bind whatever must live on the thread that will run the batches. Called - // once before any other call, on that thread. - // - // false = the executor cannot run. - virtual bool start() { - return true; - } - // A session to route tasks to. nullopt = at capacity. Every successful id // must be unique for the lifetime of the consuming Runner, even after close. virtual std::optional open_session() = 0; diff --git a/extension/llm/batching/runner.cpp b/extension/llm/batching/runner.cpp index 807337ac7d6..7a21c3910bd 100644 --- a/extension/llm/batching/runner.cpp +++ b/extension/llm/batching/runner.cpp @@ -454,11 +454,6 @@ void RunnerImpl::notify_engine_() { // --- engine thread --------------------------------------------------------- void RunnerImpl::run_() { - // Start before anything is admitted, and on this thread. - if (!executor_.start()) { - lifecycle_.store(Lifecycle::Stopping, std::memory_order_release); - } - while (is_running_()) { process_pending_commands_(); reap_cancelled_();