From 39b07a58d3ef8a91e35347f738b3af13dd88ceef Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Fri, 13 Nov 2026 13:50:16 -0800 Subject: [PATCH] Add a batching executor for native-backend packages NativeExecutor, in backends/native/batching, serves a .ptn exported with --offgraph_cache --logits_to_keep selected behind extension/llm/batching: the native counterpart of ModuleExecutor, as CudaExecutor is for CUDA. It reads the cache geometry and context length from the package, builds the engine's batched cache with room for every resident session's context, and serves sessions as the cache's sequences through SessionTable. A batch puts several sessions in one forward; each forward declares its sequences and selects the logits row each input wants; a clone shares its source's cells. The engine is the caller's: the config takes a host factory and the CacheFactory backend id of that engine's caches, as the native LLM runner does, so a model's serving code picks ggml (or another engine with a batched cache) and adds only its own preparation. native_batching_test, given a package from test/export_toy_llm.py, checks that sessions decoded together in one batch get the tokens they get alone and that a clone continues as its source; it skips without one. --- backends/native/CMakeLists.txt | 5 + backends/native/batching/BUCK | 8 + backends/native/batching/CMakeLists.txt | 54 ++++ backends/native/batching/__init__.py | 0 backends/native/batching/native_executor.cpp | 286 ++++++++++++++++++ backends/native/batching/native_executor.h | 86 ++++++ backends/native/batching/targets.bzl | 30 ++ backends/native/batching/test/__init__.py | 0 .../native/batching/test/export_toy_llm.py | 53 ++++ .../batching/test/test_native_executor.cpp | 198 ++++++++++++ 10 files changed, 720 insertions(+) create mode 100644 backends/native/batching/BUCK create mode 100644 backends/native/batching/CMakeLists.txt create mode 100644 backends/native/batching/__init__.py create mode 100644 backends/native/batching/native_executor.cpp create mode 100644 backends/native/batching/native_executor.h create mode 100644 backends/native/batching/targets.bzl create mode 100644 backends/native/batching/test/__init__.py create mode 100644 backends/native/batching/test/export_toy_llm.py create mode 100644 backends/native/batching/test/test_native_executor.cpp diff --git a/backends/native/CMakeLists.txt b/backends/native/CMakeLists.txt index 1a2911aab89..2da96e95853 100644 --- a/backends/native/CMakeLists.txt +++ b/backends/native/CMakeLists.txt @@ -35,6 +35,11 @@ if(EXECUTORCH_BUILD_NATIVE_GGML) add_subdirectory(runtime/engine/ggml) endif() +# The batching executor implements extension/llm/batching's executor seam. +if(EXECUTORCH_BUILD_NATIVE AND TARGET extension_llm_batching_util) + add_subdirectory(batching) +endif() + install( DIRECTORY runtime/ DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/executorch/backends/native/runtime diff --git a/backends/native/batching/BUCK b/backends/native/batching/BUCK new file mode 100644 index 00000000000..f8c1db3fa44 --- /dev/null +++ b/backends/native/batching/BUCK @@ -0,0 +1,8 @@ +# Any targets that should be shared between fbcode and xplat must be defined in +# targets.bzl. + +load(":targets.bzl", "define_common_targets") + +oncall("executorch") + +define_common_targets() diff --git a/backends/native/batching/CMakeLists.txt b/backends/native/batching/CMakeLists.txt new file mode 100644 index 00000000000..fe7d7cf1f1f --- /dev/null +++ b/backends/native/batching/CMakeLists.txt @@ -0,0 +1,54 @@ +# 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. + +# The batching executor for native-backend packages: the native counterpart of +# extension_llm_batching_module. Added from backends/native after +# extension/llm/batching, whose executor seam it implements. + +if(NOT EXECUTORCH_ROOT) + set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../../..) +endif() + +add_library(native_batching native_executor.cpp) +target_include_directories( + native_batching PUBLIC $ + $ +) +target_compile_options( + native_batching PRIVATE $<$:/EHsc> + $<$>:-fexceptions> +) +target_link_libraries( + native_batching + PUBLIC native_engine extension_llm_batching extension_llm_batching_util + extension_llm_cache extension_tensor +) + +install( + TARGETS native_batching + EXPORT ExecuTorchTargets + DESTINATION ${CMAKE_INSTALL_LIBDIR} + INCLUDES + DESTINATION ${_common_include_directories} +) +install( + FILES native_executor.h + DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/executorch/backends/native/batching +) + +if(BUILD_TESTING AND TARGET native_ggml_engine) + include(${EXECUTORCH_ROOT}/tools/cmake/Test.cmake) + + # Skips unless ET_NATIVE_BATCHING_PTN names the output of + # test/export_toy_llm.py. + et_cxx_test( + native_batching_test SOURCES test/test_native_executor.cpp EXTRA_LIBS + native_batching native_ggml_engine + ) + target_compile_options( + native_batching_test PRIVATE $<$>:-fexceptions> + ) +endif() diff --git a/backends/native/batching/__init__.py b/backends/native/batching/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/backends/native/batching/native_executor.cpp b/backends/native/batching/native_executor.cpp new file mode 100644 index 00000000000..5882a45e1cc --- /dev/null +++ b/backends/native/batching/native_executor.cpp @@ -0,0 +1,286 @@ +// 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 +#include +#include + +namespace ptn { + +namespace batching = executorch::extension::llm::batching; +namespace cache = executorch::extension::llm::cache; +namespace llm = executorch::extension::llm; +using executorch::runtime::Error; +using executorch::runtime::Result; + +namespace { +// get_logits_to_keep_mode's value for selected logits. +constexpr int64_t kSelectedLogits = 2; +} // namespace + +struct NativeExecutor::State { + std::string method; + std::shared_ptr package; + std::shared_ptr program; + std::unique_ptr host; + std::shared_ptr cache; + std::unique_ptr guard; + cache::BatchControl* control = nullptr; + std::unique_ptr sessions; + // Declared after the host and cache, so they outlive it. + std::unique_ptr context; + std::unique_ptr executable; + size_t logits_output = 0; + int max_step_tokens = 0; + int max_session_tokens = 0; + int32_t vocab_size = 0; +}; + +NativeExecutor::NativeExecutor(std::unique_ptr state) + : state_(std::move(state)) {} + +NativeExecutor::~NativeExecutor() = default; + +Result> NativeExecutor::create( + NativeExecutorConfig config) { + if (!config.make_host || config.max_sessions <= 0 || + config.max_session_tokens < 0) { + ET_LOG(Error, "NativeExecutor: invalid engine or session limits"); + return Error::InvalidArgument; + } + try { + auto state = std::make_unique(); + state->method = config.method; + state->package = + std::make_shared(Package::load(config.model_path)); + const ByteSpan bytes = state->package->program_bytes(); + state->program = std::make_shared( + Program::load(bytes.data(), bytes.size())); + const Method& method = state->program->get_method(config.method); + const auto geometry = offgraph_cache_geometry(method, *state->package); + if (!geometry || method.graph.input_ids.size() != 3 || + state->package->metadata_int(llm::kLogitsToKeepMode) != + kSelectedLogits) { + ET_LOG( + Error, + "NativeExecutor: %s needs an off-graph KV cache and selected " + "logits (export_hf_llm.py --offgraph_cache --logits_to_keep " + "selected)", + config.model_path.c_str()); + return Error::InvalidProgram; + } + const std::optional context = + state->package->metadata_int(llm::kMaxContextLen); + if (!context || *context <= 0 || + *context > std::numeric_limits::max() || + config.max_session_tokens > *context) { + ET_LOG( + Error, + "NativeExecutor: the session token limit must fit the package's " + "context length"); + return Error::InvalidArgument; + } + state->max_session_tokens = config.max_session_tokens > 0 + ? config.max_session_tokens + : static_cast(*context); + if (config.max_sessions > + std::numeric_limits::max() / state->max_session_tokens) { + ET_LOG(Error, "NativeExecutor: the cell pool exceeds int range"); + return Error::InvalidArgument; + } + state->max_step_tokens = max_step_tokens(method); + for (size_t i = 0; i < method.output_specs.size(); ++i) { + if (method.output_specs[i].kind == OutputKind::UserOutput) { + state->logits_output = i; + const auto& sizes = + method.graph.value(method.graph.output_ids[i]).tensor_meta().sizes; + state->vocab_size = + sizes.empty() ? 0 : static_cast(sizes.back()); + break; + } + } + if (state->vocab_size <= 0) { + ET_LOG(Error, "NativeExecutor: the method returns no logits"); + return Error::InvalidProgram; + } + + // Every resident session may fill its context, so the pool holds them + // all; storage grows with the cells in use. + auto built = cache::CacheFactory::global().build( + config.cache_backend_id, + cache::kind::kBatched, + *geometry, + cache::CacheConfig{ + .capacity = config.max_sessions * state->max_session_tokens, + .kv_dtype = static_cast(config.kv_dtype)}); + if (!built.ok()) { + ET_LOG( + Error, + "NativeExecutor: no batched cache for backend %s", + config.cache_backend_id.c_str()); + return built.error(); + } + state->cache = built.get(); + state->control = state->cache->as(); + if (state->control == nullptr) { + return Error::InvalidType; + } + const std::optional seq_limit = state->control->max_seqs(); + if (seq_limit && config.max_sessions > *seq_limit) { + ET_LOG( + Error, + "NativeExecutor: %d resident sessions requested, but the cache " + "holds %d", + config.max_sessions, + *seq_limit); + return Error::InvalidArgument; + } + state->sessions = std::make_unique( + *state->control, + config.max_sessions, + state->max_session_tokens, + state->vocab_size); + + state->guard = std::make_unique(state->cache); + executorch::runtime::BackendOptions<1> options; + const char* key = nullptr; + if (state->guard->set_option(options) != Error::Ok || + options.get_option(cache::kCacheKeyOption, key) != Error::Ok) { + return Error::Internal; + } + state->host = config.make_host(); + state->context = state->host->create_context_with_options( + state->program, state->package, {{cache::kCacheKeyOption, key}}); + return std::unique_ptr( + new NativeExecutor(std::move(state))); + } catch (const std::exception& error) { + ET_LOG(Error, "NativeExecutor: creation failed: %s", error.what()); + return Error::InvalidProgram; + } +} + +int32_t NativeExecutor::vocab_size() const { + return state_->vocab_size; +} + +int NativeExecutor::max_session_tokens() const { + return state_->max_session_tokens; +} + +bool NativeExecutor::initialize() { + try { + state_->executable = state_->context->compile(state_->method); + return true; + } catch (const std::exception& error) { + ET_LOG(Error, "NativeExecutor: compilation failed: %s", error.what()); + return false; + } +} + +size_t NativeExecutor::preferred_batch_tokens() const { + return static_cast(state_->max_step_tokens); +} + +std::optional NativeExecutor::open_session() { + return state_->sessions->open(); +} + +void NativeExecutor::close_session(batching::SessionId session) { + state_->sessions->close(session); +} + +std::optional NativeExecutor::clone( + batching::SessionId source, + batching::Position upto) { + return state_->sessions->clone(source, upto); +} + +void NativeExecutor::set_sampling( + batching::SessionId session, + const batching::SamplingParams& params, + std::optional seed) { + state_->sessions->set_sampling(session, params, seed); +} + +bool NativeExecutor::execute( + const batching::BatchInput& batch, + batching::BatchOutput& out) { + out.outputs.clear(); + out.outputs.resize(batch.inputs.size()); + try { + const auto step = state_->sessions->pack(batch); + if (!step.ok()) { + return false; + } + EngineExecutable& executable = *state_->executable; + // A batch wider than the method runs as several forwards, in order, so a + // slice attends the cells its predecessors wrote. + const int total = static_cast(step->tokens.size()); + std::vector logits; + for (int off = 0; off < total; off += state_->max_step_tokens) { + const int n = std::min(state_->max_step_tokens, total - off); + if (!state_->control->declare_step(std::vector( + step->seq_ids.begin() + off, step->seq_ids.begin() + off + n))) { + ET_LOG(Error, "NativeExecutor: the cache refused %d tokens", n); + return false; + } + const batching::util::SliceRows rows = + batching::util::select_rows(*step, off, n); + const auto k = static_cast(rows.selector.size()); + executable.resize_input(0, {1, n}); + executable.set_input(0, step->tokens.data() + off, n, ScalarType::Long); + executable.resize_input(1, {n}); + executable.set_input( + 1, step->positions.data() + off, n, ScalarType::Long); + executable.resize_input(2, {k}); + executable.set_input( + 2, rows.selector.data(), rows.selector.size(), ScalarType::Long); + executable.execute(); + if (rows.inputs.empty()) { + continue; + } + logits.resize(static_cast(k * state_->vocab_size)); + executable.get_output( + state_->logits_output, + logits.data(), + logits.size(), + ScalarType::Float); + // Non-const: the sampler reduces each row in place. + auto selected = executorch::extension::make_tensor_ptr( + {static_cast(k), state_->vocab_size}, logits); + for (size_t row = 0; row < rows.inputs.size(); ++row) { + const size_t input = rows.inputs[row]; + const batching::SessionId session = batch.inputs[input].sid; + const auto token = + state_->sessions->sample(session, *selected, static_cast(row)); + if (!token) { + return false; + } + out.outputs[input] = batching::Output{session, {*token}}; + } + } + return true; + } catch (const std::exception& error) { + ET_LOG(Error, "NativeExecutor: execution failed: %s", error.what()); + return false; + } +} + +} // namespace ptn diff --git a/backends/native/batching/native_executor.h b/backends/native/batching/native_executor.h new file mode 100644 index 00000000000..ea9a7c62933 --- /dev/null +++ b/backends/native/batching/native_executor.h @@ -0,0 +1,86 @@ +// 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 + +// A batching Executor for a native-backend .ptn: the native counterpart of +// extension/llm/batching's ModuleExecutor. The package must be exported with +// an off-graph KV cache and selected logits (export_hf_llm.py --offgraph_cache +// --logits_to_keep selected), so its method takes tokens [1, T], positions [T] +// and the token indices [K] whose logits [1, K, vocab] it returns. A session is +// one sequence of the engine's batched KV cache, a batch is one or more +// forwards carrying every input's tokens on a single axis, and a clone shares +// its source's cells. + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +namespace ptn { + +struct NativeExecutorConfig { + std::string model_path; + // Creates the engine host that runs the method. + std::function()> make_host; + // CacheFactory backend id of that engine's caches, whose kind::kBatched + // builder must be registered, e.g. kGgmlCacheBackendId. + std::string cache_backend_id; + int max_sessions = 1; + // Zero uses the package's context length. + int max_session_tokens = 0; + ScalarType kv_dtype = ScalarType::Float; + std::string method = "forward"; +}; + +class NativeExecutor final + : public executorch::extension::llm::batching::Executor { + public: + // Loads the package and builds a cache holding every session's context; + // its storage grows with the cells in use. Returns an error for a package + // without an off-graph cache or selected logits, unusable limits, or a + // cache backend without a batched builder. + static executorch::runtime::Result> create( + NativeExecutorConfig config); + ~NativeExecutor() override; + + int32_t vocab_size() const; + int max_session_tokens() const; + + // Compiles the method on the engine thread. + bool initialize() override; + // The widest forward the method takes; a wider batch runs as several. + size_t preferred_batch_tokens() const override; + std::optional open_session() + override; + void close_session( + executorch::extension::llm::batching::SessionId session) override; + std::optional clone( + executorch::extension::llm::batching::SessionId source, + executorch::extension::llm::batching::Position upto) override; + void set_sampling( + executorch::extension::llm::batching::SessionId session, + const executorch::extension::llm::batching::SamplingParams& params, + std::optional seed) override; + bool execute( + const executorch::extension::llm::batching::BatchInput& batch, + executorch::extension::llm::batching::BatchOutput& out) override; + + private: + struct State; + explicit NativeExecutor(std::unique_ptr state); + + std::unique_ptr state_; +}; + +} // namespace ptn diff --git a/backends/native/batching/targets.bzl b/backends/native/batching/targets.bzl new file mode 100644 index 00000000000..3794208c814 --- /dev/null +++ b/backends/native/batching/targets.bzl @@ -0,0 +1,30 @@ +load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime") + +def define_common_targets(): + runtime.cxx_library( + name = "native_executor", + srcs = [ + "native_executor.cpp", + ], + exported_headers = [ + "native_executor.h", + ], + visibility = ["PUBLIC"], + exported_deps = [ + "//executorch/backends/native/runtime/engine:engine", + "//executorch/backends/native/runtime/graph:scalar_type", + "//executorch/extension/llm/batching:batching", + "//executorch/runtime/core:core", + ], + deps = [ + "//executorch/backends/native/runtime:offgraph_cache", + "//executorch/backends/native/runtime:runtime", + "//executorch/backends/native/runtime/deserialize:package", + "//executorch/extension/llm/batching:session_table", + "//executorch/extension/llm/cache:kv_cache", + "//executorch/extension/llm/runner:constants", + "//executorch/extension/tensor:tensor", + "//executorch/runtime/backend:interface", + "//executorch/runtime/platform:platform", + ], + ) diff --git a/backends/native/batching/test/__init__.py b/backends/native/batching/test/__init__.py new file mode 100644 index 00000000000..e69de29bb2d diff --git a/backends/native/batching/test/export_toy_llm.py b/backends/native/batching/test/export_toy_llm.py new file mode 100644 index 00000000000..dd77a4f03c7 --- /dev/null +++ b/backends/native/batching/test/export_toy_llm.py @@ -0,0 +1,53 @@ +# 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. + +"""Export a tiny random Llama with an off-graph KV cache and selected logits, +for test_native_executor.cpp: + + python -m executorch.backends.native.batching.test.export_toy_llm toy.ptn + ET_NATIVE_BATCHING_PTN=toy.ptn native_batching_test +""" + +import sys +import tempfile +from unittest import mock + +from executorch.backends.native.examples.hf_llm import export_hf_llm +from transformers import LlamaConfig + + +def main() -> None: + output = sys.argv[1] + with tempfile.TemporaryDirectory() as model: + LlamaConfig( + vocab_size=128, + hidden_size=64, + intermediate_size=128, + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + head_dim=16, + max_position_embeddings=256, + ).save_pretrained(model) + argv = [ + "export_hf_llm", + "--model", + model, + "--random_weights", + "--offgraph_cache", + "--logits_to_keep", + "selected", + "--max_seq_length", + "16", + "--output", + output, + ] + with mock.patch.object(sys, "argv", argv): + export_hf_llm.main() + + +if __name__ == "__main__": + main() diff --git a/backends/native/batching/test/test_native_executor.cpp b/backends/native/batching/test/test_native_executor.cpp new file mode 100644 index 00000000000..34afd01691d --- /dev/null +++ b/backends/native/batching/test/test_native_executor.cpp @@ -0,0 +1,198 @@ +// 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 + +namespace { + +namespace batching = executorch::extension::llm::batching; +using batching::Token; + +class NativeExecutorTest : public ::testing::Test { + protected: + void SetUp() override { + executorch::runtime::runtime_init(); + const char* path = std::getenv("ET_NATIVE_BATCHING_PTN"); + if (path == nullptr) { + GTEST_SKIP() << "ET_NATIVE_BATCHING_PTN names no package; see " + "test/export_toy_llm.py"; + } + ptn::register_ggml_cache_builders(); + auto created = ptn::NativeExecutor::create( + {.model_path = path, + .make_host = [] { return ptn::GgmlEngineHost::create(); }, + .cache_backend_id = ptn::kGgmlCacheBackendId, + .max_sessions = 4, + .max_session_tokens = 64}); + ASSERT_TRUE(created.ok()); + executor_ = std::move(created.get()); + ASSERT_TRUE(executor_->initialize()); + } + + batching::SessionId open() { + const auto session = executor_->open_session(); + EXPECT_TRUE(session.has_value()); + executor_->set_sampling(*session, batching::SamplingParams{}, 0); + return *session; + } + + static batching::Input input( + batching::SessionId session, + std::vector tokens, + batching::Position position) { + batching::Input in; + in.sid = session; + in.size = tokens.size(); + in.offset = 0; + in.payload = batching::TokenInputPtr( + std::make_shared>(std::move(tokens))); + in.position = position; + in.produce_output = true; + return in; + } + + // Runs one batch and returns each input's next token. + std::vector step(std::vector inputs) { + batching::BatchOutput out; + EXPECT_TRUE(executor_->execute(batching::BatchInput{inputs}, out)); + std::vector tokens; + for (const auto& output : out.outputs) { + EXPECT_TRUE(output.has_value()); + tokens.push_back(output ? output->tokens.at(0) : 0); + } + return tokens; + } + + // Greedy tokens after `prompt`, decoding `steps` more, in a session alone. + std::vector generate_alone( + const std::vector& prompt, + int steps) { + const batching::SessionId session = open(); + std::vector tokens = step({input(session, prompt, 0)}); + for (int i = 1; i < steps; ++i) { + tokens.push_back(step({input( + session, + {tokens.back()}, + static_cast(prompt.size()) + i - 1)})[0]); + } + executor_->close_session(session); + return tokens; + } + + std::unique_ptr executor_; +}; + +TEST(NativeExecutorCreateTest, RefusesAMissingEngine) { + EXPECT_EQ( + ptn::NativeExecutor::create({.model_path = "unused.ptn"}).error(), + executorch::runtime::Error::InvalidArgument); +} + +// Two sessions decoded in one batch give each the tokens it gets alone. +TEST_F(NativeExecutorTest, BatchedSessionsMatchSessionsAlone) { + const std::vector a = {3, 14, 15, 92, 65}; + const std::vector b = {27, 18, 28}; + const std::vector want_a = generate_alone(a, 6); + const std::vector want_b = generate_alone(b, 6); + + const batching::SessionId sa = open(); + const batching::SessionId sb = open(); + std::vector got_a, got_b; + std::vector next = step({input(sa, a, 0), input(sb, b, 0)}); + for (int i = 0; i < 6; ++i) { + got_a.push_back(next[0]); + got_b.push_back(next[1]); + if (i + 1 < 6) { + next = step( + {input(sa, {next[0]}, static_cast(a.size()) + i), + input( + sb, {next[1]}, static_cast(b.size()) + i)}); + } + } + EXPECT_EQ(got_a, want_a); + EXPECT_EQ(got_b, want_b); +} + +// A clone continues exactly as its source does, and a prompt wider than the +// method runs as several forwards. +TEST_F(NativeExecutorTest, ClonesContinueLikeTheirSource) { + std::vector prompt; + for (Token t = 1; + t <= static_cast(executor_->preferred_batch_tokens()) + 5; + ++t) { + prompt.push_back(t % 100); + } + const auto position = static_cast(prompt.size()); + const batching::SessionId source = open(); + const Token first = step({input(source, prompt, 0)})[0]; + const auto clone = executor_->clone(source, position); + ASSERT_TRUE(clone.has_value()); + executor_->set_sampling(*clone, batching::SamplingParams{}, 0); + const std::vector next = step( + {input(source, {first}, position), input(*clone, {first}, position)}); + EXPECT_EQ(next[0], next[1]); +} + +// One batch carries a session's decode and another's prompt, wider than the +// method, so the prompt's last token is selected in a later forward. +TEST_F(NativeExecutorTest, MixedPrefillAndDecodeMatchSessionsAlone) { + const std::vector a = {3, 14, 15, 92, 65}; + std::vector b; + for (Token t = 1; + t <= static_cast(executor_->preferred_batch_tokens()) + 5; + ++t) { + b.push_back(t % 100); + } + const std::vector want_a = generate_alone(a, 4); + const std::vector want_b = generate_alone(b, 3); + + const batching::SessionId sa = open(); + const batching::SessionId sb = open(); + const auto la = static_cast(a.size()); + const auto lb = static_cast(b.size()); + std::vector got_a = step({input(sa, a, 0)}); + std::vector next = + step({input(sa, {got_a.back()}, la), input(sb, b, 0)}); + got_a.push_back(next[0]); + std::vector got_b = {next[1]}; + for (batching::Position i = 1; i < 3; ++i) { + next = step( + {input(sa, {got_a.back()}, la + i), + input(sb, {got_b.back()}, lb + i - 1)}); + got_a.push_back(next[0]); + got_b.push_back(next[1]); + } + EXPECT_EQ(got_a, want_a); + EXPECT_EQ(got_b, want_b); +} + +// A batch the cache refuses, here a position past the session's length, fails +// without changing the session, which then continues as it would have. +TEST_F(NativeExecutorTest, RefusedBatchLeavesTheSession) { + const std::vector prompt = {3, 14, 15, 92, 65}; + const std::vector want = generate_alone(prompt, 3); + const batching::SessionId session = open(); + const auto length = static_cast(prompt.size()); + std::vector got = step({input(session, prompt, 0)}); + batching::BatchOutput out; + EXPECT_FALSE(executor_->execute( + batching::BatchInput{{input(session, {got.back()}, length + 1)}}, out)); + for (batching::Position i = 0; i < 2; ++i) { + got.push_back(step({input(session, {got.back()}, length + i)})[0]); + } + EXPECT_EQ(got, want); +} + +} // namespace