From 703d73ee4f0597312e40375500eb099c4232ee5e Mon Sep 17 00:00:00 2001 From: Scott Roy Date: Wed, 11 Nov 2026 14:50:23 -0800 Subject: [PATCH] Serve one batching session from a single-sequence cache SessionTable maps sessions to cache sequences through BatchControl, which only multi-sequence caches offer, so an executor over a single-sequence cache had to reimplement packing, rewinds and sampling. SequenceBatchControl gives such a cache the BatchControl face with one sequence: seq_new claims it, seq_rm clears and frees it, rewind and pos go to SequenceControl, and clone is refused. A SessionTable over it serves one session at a time. --- extension/llm/batching/targets.bzl | 1 + .../llm/batching/test/session_table_test.cpp | 62 ++++++++++++ .../batching/util/sequence_batch_control.h | 96 +++++++++++++++++++ 3 files changed, 159 insertions(+) create mode 100644 extension/llm/batching/util/sequence_batch_control.h diff --git a/extension/llm/batching/targets.bzl b/extension/llm/batching/targets.bzl index b0aa4606f51..90f033ee099 100644 --- a/extension/llm/batching/targets.bzl +++ b/extension/llm/batching/targets.bzl @@ -51,6 +51,7 @@ def define_common_targets(): "util/session_table.cpp", ], exported_headers = [ + "util/sequence_batch_control.h", "util/session_table.h", ], visibility = ["PUBLIC"], diff --git a/extension/llm/batching/test/session_table_test.cpp b/extension/llm/batching/test/session_table_test.cpp index 0efcb0a36fa..ad828104c4a 100644 --- a/extension/llm/batching/test/session_table_test.cpp +++ b/extension/llm/batching/test/session_table_test.cpp @@ -6,9 +6,11 @@ * LICENSE file in the root directory of this source tree. */ +#include #include #include +#include #include #include @@ -251,3 +253,63 @@ TEST(SelectRowsTest, KeepsOnlyTheSlicesWantedRows) { EXPECT_EQ(rows.selector, std::vector({2})); EXPECT_TRUE(rows.inputs.empty()); } + +// One session over a single-sequence cache, through the same table. +TEST(SequenceBatchControlTest, ServesOneSessionAtATime) { + ::executorch::runtime::runtime_init(); + cache::CacheGeometry geometry; + geometry.layers = {{{cache::LayerPolicy::Kind::Flat, 0}, 1, 1}}; + cache::SequenceCache sequence(geometry, cache::CacheConfig{16, 6}); + util::SequenceBatchControl control(sequence); + util::SessionTable table(control, /*max_sessions=*/2, 16, kVocab); + // Commits a step as the cache's own byte layer would. + const auto commit = [&](const util::PackedStep& step) { + ASSERT_TRUE(control.declare_step(step.seq_ids)); + const auto plan = sequence.plan( + 0, + static_cast(step.positions.front()), + static_cast(step.positions.size())); + ASSERT_TRUE(plan.has_value()); + sequence.commit(*plan); + }; + const auto input = [](batching::SessionId sid, + std::vector tokens, + batching::Position position) { + batching::Input in; + in.sid = sid; + 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; + }; + + const auto session = table.open(); + ASSERT_TRUE(session.has_value()); + EXPECT_FALSE(table.open().has_value()); + EXPECT_FALSE(table.clone(*session, 1).has_value()); + + auto step = table.pack({{input(*session, {1, 2, 3}, 0)}}); + ASSERT_TRUE(step.ok()); + EXPECT_EQ(step->seq_ids, (std::vector{0, 0, 0})); + EXPECT_EQ(step->positions, (std::vector{0, 1, 2})); + commit(*step); + EXPECT_EQ(control.pos(0), 3); + + // A step starting inside the committed history rewinds to its start. + const auto rewound = table.pack({{input(*session, {7}, 1)}}); + ASSERT_TRUE(rewound.ok()); + EXPECT_EQ(control.pos(0), 1); + commit(*rewound); + EXPECT_EQ(control.pos(0), 2); + EXPECT_FALSE(table.pack({{input(*session, {8}, 4)}}).ok()); + + table.close(*session); + EXPECT_EQ(sequence.length(), 0); + const auto reopened = table.open(); + ASSERT_TRUE(reopened.has_value()); + EXPECT_NE(*reopened, *session); + EXPECT_EQ(control.pos(0), 0); +} diff --git a/extension/llm/batching/util/sequence_batch_control.h b/extension/llm/batching/util/sequence_batch_control.h new file mode 100644 index 00000000000..4ccb30dd875 --- /dev/null +++ b/extension/llm/batching/util/sequence_batch_control.h @@ -0,0 +1,96 @@ +/* + * 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 + +// BatchControl over a cache that holds one sequence, so a SessionTable serves +// one session from it: an executor over a single-sequence cache. Sequence 0 is +// the only id, live from seq_new until seq_rm, which clears it; clone is +// refused, since the cache has nowhere to put a second history. + +#include +#include +#include +#include + +#include +#include // ET_EXPERIMENTAL + +namespace executorch { +namespace extension { +namespace llm { +namespace batching { +namespace util { + +class ET_EXPERIMENTAL SequenceBatchControl final + : public ::executorch::extension::llm::cache::BatchControl { + public: + // `sequence` must outlive this. + explicit SequenceBatchControl( + ::executorch::extension::llm::cache::SequenceControl& sequence) + : sequence_(sequence) {} + + int capacity() const override { + return sequence_.capacity(); + } + void clear() override { + sequence_.clear(); + } + bool declare_step(const std::vector& seq_ids) override { + return live_ && + std::all_of( + seq_ids.begin(), + seq_ids.end(), + [](std::int32_t id) { return id == 0; }) && + sequence_.can_extend(static_cast(seq_ids.size())); + } + std::optional max_seqs() const override { + return 1; + } + std::optional seq_new() override { + if (live_) { + return std::nullopt; + } + sequence_.clear(); + live_ = true; + return 0; + } + std::optional seq_clone( + std::int32_t /*src*/, + std::optional /*upto*/) override { + return std::nullopt; + } + bool seq_rm(std::int32_t seq_id) override { + if (!holds(seq_id)) { + return false; + } + sequence_.clear(); + live_ = false; + return true; + } + bool rewind(std::int32_t seq_id, int position) override { + return holds(seq_id) && sequence_.rewind(position); + } + int pos(std::int32_t seq_id) const override { + return holds(seq_id) ? sequence_.length() : 0; + } + + private: + bool holds(std::int32_t seq_id) const { + return live_ && seq_id == 0; + } + + ::executorch::extension::llm::cache::SequenceControl& sequence_; + bool live_ = false; +}; + +} // namespace util +} // namespace batching +} // namespace llm +} // namespace extension +} // namespace executorch