Serve one batching session from a single-sequence cache

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"],

extension/llm/batching/test/session_table_test.cpp

@@ -6,9 +6,11 @@
  * LICENSE file in the root directory of this source tree.
  */
 
+#include <executorch/extension/llm/batching/util/sequence_batch_control.h>
 #include <executorch/extension/llm/batching/util/session_table.h>
 
 #include <executorch/extension/llm/cache/cell_cache.h>
+#include <executorch/extension/llm/cache/sequence_cache.h>
 #include <executorch/extension/tensor/tensor.h>
 #include <executorch/runtime/platform/runtime.h>
 
@@ -251,3 +253,63 @@ TEST(SelectRowsTest, KeepsOnlyTheSlicesWantedRows) {
   EXPECT_EQ(rows.selector, std::vector<std::int64_t>({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<int>(step.positions.front()),
+        static_cast<int>(step.positions.size()));
+    ASSERT_TRUE(plan.has_value());
+    sequence.commit(*plan);
+  };
+  const auto input = [](batching::SessionId sid,
+                        std::vector<batching::Token> 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::vector<batching::Token>>(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<int32_t>{0, 0, 0}));
+  EXPECT_EQ(step->positions, (std::vector<int64_t>{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);
+}

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 <algorithm>
+#include <cstdint>
+#include <optional>
+#include <vector>
+
+#include <executorch/extension/llm/cache/cache.h>
+#include <executorch/runtime/platform/compiler.h> // 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<std::int32_t>& 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<int>(seq_ids.size()));
+  }
+  std::optional<int> max_seqs() const override {
+    return 1;
+  }
+  std::optional<std::int32_t> seq_new() override {
+    if (live_) {
+      return std::nullopt;
+    }
+    sequence_.clear();
+    live_ = true;
+    return 0;
+  }
+  std::optional<std::int32_t> seq_clone(
+      std::int32_t /*src*/,
+      std::optional<int> /*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