diff --git a/CMakeLists.txt b/CMakeLists.txt index 40cbef38abe..ef7e9479e01 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -996,6 +996,8 @@ if(EXECUTORCH_BUILD_EXTENSION_LLM) list(APPEND _executorch_extensions tokenizers) 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) endif() if(EXECUTORCH_BUILD_EXTENSION_RUNNER_UTIL) diff --git a/extension/llm/batching/CMakeLists.txt b/extension/llm/batching/CMakeLists.txt new file mode 100644 index 00000000000..548593d1f4d --- /dev/null +++ b/extension/llm/batching/CMakeLists.txt @@ -0,0 +1,43 @@ +# 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. + +# Batched LLM serving: the step scheduler, and the runner that drives batches +# through a model. Currently header-only and free of ExecuTorch runtime types, +# so it is an INTERFACE target; this becomes a real library once the runner +# lands with a .cpp. + +if(NOT EXECUTORCH_ROOT) + set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../../..) +endif() + +add_library(extension_llm_batching INTERFACE) +# std::optional in the public headers. +target_compile_features(extension_llm_batching INTERFACE cxx_std_17) +target_include_directories( + extension_llm_batching INTERFACE ${_common_include_directories} +) +target_compile_options( + extension_llm_batching INTERFACE ${_common_compile_options} +) + +install( + TARGETS extension_llm_batching + EXPORT ExecuTorchTargets + DESTINATION ${CMAKE_INSTALL_LIBDIR} + INCLUDES + DESTINATION ${_common_include_directories} +) +install( + DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/ + DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/executorch/extension/llm/batching + FILES_MATCHING + PATTERN "*.h" + PATTERN "test" EXCLUDE +) + +if(BUILD_TESTING) + add_subdirectory(test) +endif() diff --git a/extension/llm/batching/decode_first_scheduler.h b/extension/llm/batching/decode_first_scheduler.h new file mode 100644 index 00000000000..62f4d0b536d --- /dev/null +++ b/extension/llm/batching/decode_first_scheduler.h @@ -0,0 +1,375 @@ +/* + * 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 + +// Decodes first, up to max_decode_sequences, then spends the rest of +// max_batch_tokens on prefill, so prefill never delays a queued decode. +// Prefill rotates over sessions taking one chunk each, so a long prompt +// cannot monopolise a batch. +// +// Reads only input.size, input.sid, is_decode and tid from a Task. The rest is +// carried through untouched. + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include // ET_EXPERIMENTAL + +namespace executorch { +namespace extension { +namespace llm { +namespace batching { + +// Shared so that pending_ and the queues name the same task. Cancelling marks +// it once, and each queue skips it when scheduling reaches it. +using TaskPtr = std::shared_ptr; +using TaskQueue = std::deque; + +class ET_EXPERIMENTAL DecodeFirstScheduler : public Scheduler { + public: + // Returns nullptr if the limits are unusable. All three must be non-zero, + // and the budget must cover a saturated decode batch plus one full chunk. + // Below that floor a chunk of max_prefill_chunk_size could be admitted and + // then never fit in any batch, stranding the task rather than delaying it. + // The same floor is what lets take_decodes_ spend without checking the + // budget. + // + // A factory rather than a throwing constructor, because + // EXECUTORCH_OPTIMIZE_SIZE builds with -fno-exceptions, and this header is + // deliberately free of ExecuTorch runtime types, so ET_CHECK is unavailable + // too. + static std::unique_ptr create( + std::size_t max_batch_tokens = 544, + std::size_t max_decode_sequences = 32, + std::size_t max_prefill_chunk_size = 256) { + if (max_batch_tokens == 0 || max_decode_sequences == 0 || + max_prefill_chunk_size == 0) { + return nullptr; + } + if (max_decode_sequences >= max_batch_tokens) { + return nullptr; // no room left for prefill + } + // Ordered by the check above, so the subtraction cannot wrap. + if (max_prefill_chunk_size > max_batch_tokens - max_decode_sequences) { + return nullptr; + } + return std::unique_ptr(new DecodeFirstScheduler( + max_batch_tokens, max_decode_sequences, max_prefill_chunk_size)); + } + + // Rejects an empty task, a decode wider than one token, an oversized chunk, + // and a tid already queued, including one repeated inside this vector. + bool submit(std::vector tasks) override { + std::lock_guard g(mutex_); + for (std::size_t i = 0; i < tasks.size(); ++i) { + if (!admissible_(tasks[i])) { + return false; + } + // admissible_ tests pending_, which this vector has not joined yet, so a + // repeat within one submit has to be caught here. Pairwise because a + // vector holds one prompt's chunks: small, and it allocates nothing. + for (std::size_t j = 0; j < i; ++j) { + if (tasks[j].tid == tasks[i].tid) { + return false; + } + } + } + for (Task& t : tasks) { + enqueue_(std::move(t)); + } + return true; + } + + bool has_work() const override { + std::lock_guard g(mutex_); + return !pending_.empty(); + } + + // Spending accumulates instead of counting a budget down, so no arithmetic + // here can wrap below zero. + std::vector get_work() override { + std::vector taken; + std::lock_guard g(mutex_); + std::size_t spent = 0; + // Sessions that already hold a decode in this batch. The executor is + // promised consecutive ranges and one produce_output per session, so a + // session's second decode, or a prefill chunk beside its decode, waits for + // the next batch. Local to the call, so unlike a persistent index it + // cannot fall out of sync with the queues. + std::unordered_set decoding; + + take_decodes_(taken, spent, decoding); + while (spent < max_batch_tokens_ && + take_prefill_pass_(taken, spent, decoding)) { + } + return taken; + } + + std::vector cancel(SessionId sid) override { + std::vector dropped; + std::lock_guard g(mutex_); + + auto prefills = prefill_by_session_.find(sid); + if (prefills != prefill_by_session_.end()) { + for (const TaskPtr& t : prefills->second) { + if (release_(t)) { + dropped.push_back(std::move(*t)); + } + } + // The entry stays behind, empty, because it holds the session's place in + // the rotation. take_prefill_pass_ retires the two together. + prefills->second.clear(); + } + // Marked in place rather than erased: the queues already tolerate + // cancelled entries, and drop_cancelled_ removes them as they surface. + for (const TaskPtr& t : decode_queue_) { + if (t->input.sid == sid && release_(t)) { + dropped.push_back(std::move(*t)); + } + } + return dropped; + } + + std::vector clear() override { + std::vector dropped; + std::lock_guard g(mutex_); + + dropped.reserve(pending_.size()); + for (auto& entry : pending_) { + entry.second->cancelled = true; + dropped.push_back(std::move(*entry.second)); + } + pending_.clear(); + decode_queue_.clear(); + prefill_by_session_.clear(); + prefill_rotation_.clear(); + return dropped; + } + + // The whole budget one batch may spend, in tokens. + std::size_t max_batch_tokens() const { + return max_batch_tokens_; + } + // Decodes admitted per batch. The rest wait. + std::size_t max_decode_sequences() const { + return max_decode_sequences_; + } + // Largest chunk accepted. A larger submit is rejected, not split. + std::size_t max_prefill_chunk_size() const override { + return max_prefill_chunk_size_; + } + + private: + DecodeFirstScheduler( + std::size_t max_batch_tokens, + std::size_t max_decode_sequences, + std::size_t max_prefill_chunk_size) + : max_batch_tokens_(max_batch_tokens), + max_decode_sequences_(max_decode_sequences), + max_prefill_chunk_size_(max_prefill_chunk_size) {} + + // Caller holds mutex_. + bool admissible_(const Task& t) const { + if (t.input.size == 0) { + return false; + } + if (t.is_decode) { + // take_decodes_ spends one token per decode, and create()'s floor is + // written in those terms, so a wider decode would overspend the budget. + if (t.input.size != 1) { + return false; + } + } else if (t.input.size > max_prefill_chunk_size_) { + return false; + } + return pending_.find(t.tid) == pending_.end(); + } + + // Caller holds mutex_, having already checked admissible_. + void enqueue_(Task&& task) { + const bool decode = task.is_decode; + const SessionId sid = task.input.sid; + const TaskId tid = task.tid; + + auto t = std::make_shared(std::move(task)); + t->cancelled = false; + pending_.emplace(tid, t); + + if (decode) { + decode_queue_.push_back(t); + return; + } + // A session is in the rotation exactly when it has an entry here, so the + // insertion, not the emptiness of the queue, is what decides whether it + // joins. Testing for emptiness would give a second slot to a session whose + // entry cancel() emptied, and it would then take two turns per pass. + auto entry = prefill_by_session_.try_emplace(sid); + if (entry.second) { + prefill_rotation_.push_back(sid); + } + entry.first->second.push_back(t); + } + + // pending_ tracks only tasks still waiting, so a dispatched one leaves it. + // Caller holds mutex_. + void dispatch_(TaskId tid) { + pending_.erase(tid); + } + + // Returns whether this call was the one that dropped the task, so that a + // double cancel reports it once. Caller holds mutex_. + bool release_(const TaskPtr& t) { + if (pending_.erase(t->tid) == 0) { + return false; // already handed out, or already dropped + } + t->cancelled = true; + return true; + } + + // Discards entries release_ already dropped. They are still in the queue + // only because a deque cannot erase from the middle. + static void drop_cancelled_(TaskQueue& q) { + while (!q.empty() && q.front()->cancelled) { + q.pop_front(); + } + } + + // Takes decodes in arrival order, at most one per session. create() keeps + // max_batch_tokens above max_decode_sequences, so this cannot exhaust the + // budget. Caller holds mutex_. + void take_decodes_( + std::vector& taken, + std::size_t& spent, + std::unordered_set& decoding) { + std::vector deferred; // session already decoding in this batch + std::size_t n = 0; + while (n < max_decode_sequences_) { + drop_cancelled_(decode_queue_); + if (decode_queue_.empty()) { + break; + } + const TaskPtr t = decode_queue_.front(); + decode_queue_.pop_front(); + if (!decoding.insert(t->input.sid).second) { + deferred.push_back(t); + continue; + } + // Read the id before moving the task out of the shared entry. + const TaskId tid = t->tid; + taken.push_back(std::move(*t)); + dispatch_(tid); + ++n; + spent += 1; + } + // Returned to the head in arrival order, so a session passed over here + // still leads the queue next time. + for (auto it = deferred.rbegin(); it != deferred.rend(); ++it) { + decode_queue_.push_front(*it); + } + } + + // One turn each, in rotation order. Returns whether anything was taken. + // Caller holds mutex_. + bool take_prefill_pass_( + std::vector& taken, + std::size_t& spent, + const std::unordered_set& decoding) { + bool progress = false; + std::vector deferred; // passed over, took nothing + std::vector served; + + // Nothing rejoins the rotation inside this loop, so each session is + // visited at most once per pass. + while (spent < max_batch_tokens_ && !prefill_rotation_.empty()) { + const SessionId sid = prefill_rotation_.front(); + prefill_rotation_.pop_front(); + // A rotation slot exists exactly when the map entry does. enqueue_ adds + // the pair, and every path below retires the pair. + auto it = prefill_by_session_.find(sid); + assert( + it != prefill_by_session_.end() && + "prefill_rotation_ names a session with no queue"); + TaskQueue& dq = it->second; + drop_cancelled_(dq); + if (dq.empty()) { + prefill_by_session_.erase(it); + continue; + } + if (decoding.count(sid) != 0) { + // The session's decode is already in this batch. Adding a chunk beside + // it would hand the executor two ranges that do not adjoin, both + // asking to produce output. + deferred.push_back(sid); + continue; + } + const std::size_t n = dq.front()->input.size; + // The loop condition keeps spent below the budget, so the remaining room + // is positive. + if (n > max_batch_tokens_ - spent) { + // Pass the turn on rather than stop, since a smaller chunk behind this + // one may still fit. + deferred.push_back(sid); + continue; + } + const TaskPtr t = dq.front(); + const TaskId tid = t->tid; + taken.push_back(std::move(*t)); + dispatch_(tid); + dq.pop_front(); + spent += n; + progress = true; + if (dq.empty()) { + prefill_by_session_.erase(it); + } else { + served.push_back(sid); + } + } + + // Sessions that took nothing keep their order and rank ahead of those that + // did. They were popped before whatever remains in the rotation, so + // restoring them at the head preserves arrival order. + for (auto it = deferred.rbegin(); it != deferred.rend(); ++it) { + prefill_rotation_.push_front(*it); + } + for (SessionId sid : served) { + prefill_rotation_.push_back(sid); + } + return progress; + } + + mutable std::mutex mutex_; + std::size_t max_batch_tokens_; + std::size_t max_decode_sequences_; + std::size_t max_prefill_chunk_size_; + + TaskQueue decode_queue_; + + // A session is in prefill_rotation_ exactly when it has an entry here. That + // pairing is what keeps the rotation free of duplicates. + std::unordered_map prefill_by_session_; + std::deque prefill_rotation_; + + // Membership is what "still queued" means, so no separate count can fall out + // of sync with it. + std::unordered_map pending_; +}; + +} // namespace batching +} // namespace llm +} // namespace extension +} // namespace executorch diff --git a/extension/llm/batching/scheduler.h b/extension/llm/batching/scheduler.h new file mode 100644 index 00000000000..b578a6250a1 --- /dev/null +++ b/extension/llm/batching/scheduler.h @@ -0,0 +1,73 @@ +/* + * 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 + +// Decides which tasks run in the next batch, and accepts submissions. +// +// Bookkeeping only: never runs a callback, never calls an executor. Tasks are +// handed back and the caller decides what happens to them. +// +// Implementations must be thread safe, and cheap enough for an async context: +// bounded work, no I/O, no blocking on caller code. + +#include +#include + +#include +#include // ET_EXPERIMENTAL + +namespace executorch { +namespace extension { +namespace llm { +namespace batching { + +class ET_EXPERIMENTAL Scheduler { + public: + virtual ~Scheduler() = default; + + // All or nothing: a prompt's chunks only make sense together. False means + // rejected, and nothing was queued. + // + // A tid identifies a task to get_work() and cancel(), so it must be unique + // among the tasks queued here, including the others in this vector. It is + // free for reuse once the task has been dispatched or cancelled; ids need + // not be unique for all time. + virtual bool submit(std::vector tasks) = 0; + + // A hint, not a reservation. Another thread may take the work first. + virtual bool has_work() const = 0; + + // The caller owns what it gets back and must complete each task once. + // + // A session may appear more than once, but only as consecutive prefill + // chunks, which form one wider prefill. At most one of its tasks has + // produce_output. Submitting chunks whose positions abut is the caller's + // responsibility; the scheduler preserves their order but does not check + // the positions. + virtual std::vector get_work() = 0; + + // Drops the session's queued tasks and returns them, to be completed as + // Cancelled. A task already handed out belongs to the caller, so cancelling + // an in-flight task, or an unknown session, returns nothing. + virtual std::vector cancel(SessionId sid) = 0; + + // Drops every queued task and returns them all, for shutdown. + // For shutdown. + virtual std::vector clear() = 0; + + // Largest prefill chunk this scheduler will admit. Callers split a prompt to + // fit; a wider chunk is rejected, not split. Non-zero, so it is safe to + // divide by. + virtual std::size_t max_prefill_chunk_size() const = 0; +}; + +} // namespace batching +} // namespace llm +} // namespace extension +} // namespace executorch diff --git a/extension/llm/batching/test/CMakeLists.txt b/extension/llm/batching/test/CMakeLists.txt new file mode 100644 index 00000000000..603fcf1e341 --- /dev/null +++ b/extension/llm/batching/test/CMakeLists.txt @@ -0,0 +1,18 @@ +# 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. + +cmake_minimum_required(VERSION 3.19) + +set(EXECUTORCH_ROOT ${CMAKE_CURRENT_SOURCE_DIR}/../../../..) + +include(${EXECUTORCH_ROOT}/tools/cmake/Test.cmake) + +set(_test_srcs scheduler_test.cpp) + +et_cxx_test( + extension_llm_batching_test SOURCES ${_test_srcs} EXTRA_LIBS + extension_llm_batching +) diff --git a/extension/llm/batching/test/scheduler_test.cpp b/extension/llm/batching/test/scheduler_test.cpp new file mode 100644 index 00000000000..f1a511c5ff6 --- /dev/null +++ b/extension/llm/batching/test/scheduler_test.cpp @@ -0,0 +1,940 @@ +/* + * 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 +#include +#include + +#include + +using executorch::extension::llm::batching::BatchInput; +using executorch::extension::llm::batching::DecodeFirstScheduler; +using executorch::extension::llm::batching::Input; +using executorch::extension::llm::batching::Position; +using executorch::extension::llm::batching::Scheduler; +using executorch::extension::llm::batching::SessionId; +using executorch::extension::llm::batching::Task; +using executorch::extension::llm::batching::TaskId; +using executorch::extension::llm::batching::to_batch_input; +using executorch::extension::llm::batching::Token; + +namespace { + +using SchedulerPtr = std::unique_ptr; + +// Sizes a scheduler by decode cap and chunk size, giving it room for two full +// chunks beside a saturated decode batch. Every case below states its limits +// this way, so batch sizes stay comparable across tests. Always valid: the +// floor is decodes + chunk, and 2 * chunk + decodes clears it for any +// non-zero pair. +SchedulerPtr make_scheduler( + std::size_t max_decode_sequences, + std::size_t max_prefill_chunk_size) { + return DecodeFirstScheduler::create( + 2 * max_prefill_chunk_size + max_decode_sequences, + max_decode_sequences, + max_prefill_chunk_size); +} + +// Token values are irrelevant to scheduling, so they are all the same. +Task make_task( + TaskId task_id, + SessionId session, + std::size_t n_tokens, + Position position, + bool is_decode, + bool produce_output = true) { + auto tokens = std::make_shared>(n_tokens, 7); + Task task; + task.tid = task_id; + task.cancelled = false; + task.is_decode = is_decode; + task.input = + Input{session, produce_output, 0, n_tokens, std::move(tokens), position}; + return task; +} + +// One chunk of a prompt. A chunk of one token is still prefill, which is the +// distinction is_decode exists to make. +Task prefill( + TaskId task_id, + SessionId session, + std::size_t n_tokens, + Position position, + bool produce_output = true) { + return make_task( + task_id, + session, + n_tokens, + position, + /*is_decode=*/false, + produce_output); +} + +// One decode step, always exactly one token. +Task decode(TaskId task_id, SessionId session, Position position = 0) { + return make_task( + task_id, session, /*n_tokens=*/1, position, /*is_decode=*/true); +} + +bool submit(DecodeFirstScheduler& scheduler, Task task) { + std::vector tasks; + tasks.push_back(std::move(task)); + return scheduler.submit(std::move(tasks)); +} + +std::vector ids(const std::vector& tasks) { + std::vector out; + out.reserve(tasks.size()); + for (const Task& task : tasks) { + out.push_back(task.tid); + } + return out; +} + +std::vector sorted_ids(const std::vector& tasks) { + std::vector out = ids(tasks); + std::sort(out.begin(), out.end()); + return out; +} + +std::vector sessions(const std::vector& tasks) { + std::vector out; + out.reserve(tasks.size()); + for (const Task& task : tasks) { + out.push_back(task.input.sid); + } + return out; +} + +std::size_t token_count(const std::vector& tasks) { + std::size_t count = 0; + for (const Task& task : tasks) { + count += task.input.size; + } + return count; +} + +// The executor is promised at most one produce_output per session. A session +// may still appear more than once, because several consecutive prefill chunks +// in one batch form one wider prefill, and only the last asks for output. +bool at_most_one_output_per_session(const std::vector& tasks) { + std::set producing; + for (const Task& task : tasks) { + if (task.input.produce_output && !producing.insert(task.input.sid).second) { + return false; + } + } + return true; +} + +// Stricter, for batches whose tasks are all decodes. A decode adjoins nothing, +// so a session may hold only one slot. +bool at_most_one_task_per_session(const std::vector& tasks) { + std::set seen; + for (const Task& task : tasks) { + if (!seen.insert(task.input.sid).second) { + return false; + } + } + return true; +} + +} // namespace + +static_assert( + std::is_abstract::value, + "Scheduler is an interface and tests must use an implementation"); +static_assert( + std::is_base_of::value, + "DecodeFirstScheduler must implement Scheduler"); + +// --- construction ---------------------------------------------------------- + +TEST(CreateTest, Defaults) { + SchedulerPtr scheduler = DecodeFirstScheduler::create(); + ASSERT_NE(scheduler, nullptr); + EXPECT_EQ(scheduler->max_batch_tokens(), 544u); + EXPECT_EQ(scheduler->max_decode_sequences(), 32u); + EXPECT_EQ(scheduler->max_prefill_chunk_size(), 256u); +} + +TEST(CreateTest, RejectsZeroLimits) { + EXPECT_EQ(DecodeFirstScheduler::create(0, 32, 256), nullptr); + EXPECT_EQ(DecodeFirstScheduler::create(544, 0, 256), nullptr); + EXPECT_EQ(DecodeFirstScheduler::create(544, 32, 0), nullptr); +} + +// Below this floor a full-size chunk could be admitted and then never fit in +// any batch, leaving the task queued forever rather than merely delayed. +TEST(CreateTest, RequiresRoomForDecodesPlusAFullChunk) { + EXPECT_NE(DecodeFirstScheduler::create(288, 32, 256), nullptr) + << "exactly the floor must be accepted"; + EXPECT_EQ(DecodeFirstScheduler::create(287, 32, 256), nullptr); + EXPECT_EQ(DecodeFirstScheduler::create(100, 32, 256), nullptr); +} + +TEST(CreateTest, RejectsDecodeCapThatLeavesNoRoomForPrefill) { + EXPECT_EQ(DecodeFirstScheduler::create(544, 544, 256), nullptr); + EXPECT_EQ(DecodeFirstScheduler::create(544, 600, 256), nullptr); +} + +// The limits are size_t, so an absurd chunk size is representable and must be +// caught by the floor rather than wrapping into something that passes. +TEST(CreateTest, RejectsChunkLargerThanTheBudget) { + constexpr std::size_t kHuge = std::size_t{1} << 40; + EXPECT_EQ(DecodeFirstScheduler::create(544, 32, kHuge), nullptr); + EXPECT_NE( + DecodeFirstScheduler::create( + std::numeric_limits::max(), 32, 256), + nullptr) + << "only the floor constrains the budget; spending accumulates toward it " + "and never wraps, so a large one is not itself an error"; +} + +// --- decode and prefill are distinguished by is_decode, not by size --------- + +// Inferring the kind from input.size would send a one-token final chunk to the +// decode queue, ahead of queued decodes and outside the prefill rotation. +TEST(ClassificationTest, OneTokenPrefillIsNotADecode) { + SchedulerPtr scheduler = make_scheduler(2, 4); + ASSERT_TRUE(submit(*scheduler, prefill(1, 10, 1, 0))); + ASSERT_TRUE(submit(*scheduler, decode(2, 20))); + + EXPECT_EQ(ids(scheduler->get_work()), (std::vector{2, 1})) + << "the real decode must be taken before the one-token prefill"; +} + +// take_decodes_ spends one token per decode, and the create() floor is written +// in those terms, so a wider decode would overspend the budget. +TEST(ClassificationTest, RejectsDecodeWiderThanOneToken) { + SchedulerPtr scheduler = make_scheduler(2, 4); + EXPECT_FALSE(submit(*scheduler, make_task(1, 10, 2, 0, /*is_decode=*/true))); + EXPECT_FALSE(submit(*scheduler, make_task(2, 10, 4, 0, /*is_decode=*/true))); + EXPECT_FALSE(scheduler->has_work()); + + EXPECT_TRUE(submit(*scheduler, decode(3, 10))); +} + +// --- admission ------------------------------------------------------------- + +TEST(SubmitTest, RejectsEmptyStep) { + SchedulerPtr scheduler = make_scheduler(1, 4); + EXPECT_FALSE(submit(*scheduler, prefill(1, 10, 0, 0))); + EXPECT_FALSE(scheduler->has_work()); + EXPECT_TRUE(scheduler->get_work().empty()); +} + +TEST(SubmitTest, RejectsPrefillAboveChunkSize) { + SchedulerPtr scheduler = make_scheduler(1, 4); + EXPECT_FALSE(submit(*scheduler, prefill(1, 10, 5, 0))); + EXPECT_FALSE(scheduler->has_work()); +} + +TEST(SubmitTest, RejectsDuplicateTaskIdAlreadyWaiting) { + SchedulerPtr scheduler = make_scheduler(2, 4); + EXPECT_TRUE(submit(*scheduler, decode(7, 10))); + EXPECT_FALSE(submit(*scheduler, decode(7, 20))); + + EXPECT_EQ(ids(scheduler->get_work()), (std::vector{7})); + EXPECT_FALSE(scheduler->has_work()) << "rejection must not add a task"; +} + +// admissible_ tests pending_, which the vector being submitted has not joined +// yet, so a repeat inside one submit has to be caught separately. Accepting it +// would queue a task that pending_ never recorded. +TEST(SubmitTest, RejectsTaskIdRepeatedWithinOneSubmit) { + SchedulerPtr scheduler = make_scheduler(2, 4); + std::vector tasks; + tasks.push_back(prefill(5, 10, 4, 0, /*produce_output=*/false)); + tasks.push_back(prefill(5, 10, 4, 4)); + + EXPECT_FALSE(scheduler->submit(std::move(tasks))); + EXPECT_FALSE(scheduler->has_work()); + EXPECT_TRUE(scheduler->get_work().empty()); + EXPECT_TRUE(scheduler->clear().empty()) << "neither copy may be retained"; +} + +TEST(SubmitTest, RejectedGroupEntersNoQueue) { + SchedulerPtr scheduler = make_scheduler(2, 4); + std::vector tasks; + tasks.push_back(prefill(1, 10, 4, 0)); + tasks.push_back(prefill(2, 10, 0, 4)); + tasks.push_back(prefill(3, 10, 4, 4)); + + EXPECT_FALSE(scheduler->submit(std::move(tasks))); + EXPECT_FALSE(scheduler->has_work()); + EXPECT_TRUE(scheduler->get_work().empty()); +} + +TEST(SubmitTest, DuplicatePendingIdRejectsTheWholeGroup) { + SchedulerPtr scheduler = make_scheduler(2, 4); + EXPECT_TRUE(submit(*scheduler, decode(7, 10))); + + std::vector tasks; + tasks.push_back(prefill(8, 20, 4, 0)); + tasks.push_back(prefill(7, 20, 4, 4)); + EXPECT_FALSE(scheduler->submit(std::move(tasks))); + + EXPECT_EQ(ids(scheduler->get_work()), (std::vector{7})); + EXPECT_FALSE(scheduler->has_work()) + << "task 8 must not be partially accepted"; +} + +// Ids must be unique among queued tasks, not for all time. +TEST(SubmitTest, DispatchFreesTheTaskIdForReuse) { + SchedulerPtr scheduler = make_scheduler(2, 4); + EXPECT_TRUE(submit(*scheduler, decode(7, 10))); + ASSERT_EQ(ids(scheduler->get_work()), (std::vector{7})); + + EXPECT_TRUE(submit(*scheduler, decode(7, 10, 1))); + EXPECT_EQ(ids(scheduler->get_work()), (std::vector{7})); +} + +TEST(SubmitTest, CancelFreesTheTaskIdForReuse) { + SchedulerPtr scheduler = make_scheduler(2, 4); + EXPECT_TRUE(submit(*scheduler, decode(7, 10))); + ASSERT_EQ(ids(scheduler->cancel(10)), (std::vector{7})); + + EXPECT_TRUE(submit(*scheduler, decode(7, 10, 1))); + EXPECT_EQ(ids(scheduler->get_work()), (std::vector{7})); +} + +// --- payload passthrough --------------------------------------------------- + +TEST(PayloadTest, CarriesPayloadUninspected) { + SchedulerPtr scheduler = make_scheduler(1, 8); + auto tokens = std::make_shared>( + std::initializer_list{10, 11, 12, 13, 14, 15}); + Task task = prefill(1, 40, 4, 90, /*produce_output=*/false); + task.input.tokens = tokens; + task.input.offset = 1; + EXPECT_TRUE(submit(*scheduler, std::move(task))); + + std::vector work = scheduler->get_work(); + ASSERT_EQ(work.size(), 1u); + EXPECT_EQ(work[0].input.sid, 40); + EXPECT_FALSE(work[0].input.produce_output); + EXPECT_EQ(work[0].input.offset, 1u); + EXPECT_EQ(work[0].input.size, 4u); + EXPECT_EQ(work[0].input.tokens.get(), tokens.get()); + EXPECT_EQ(work[0].input.position, 90); +} + +TEST(PayloadTest, BatchInputKeepsTaskOrderAndSlices) { + SchedulerPtr scheduler = make_scheduler(2, 8); + EXPECT_TRUE(submit(*scheduler, decode(1, 10))); + EXPECT_TRUE(submit(*scheduler, decode(2, 20))); + EXPECT_TRUE(submit(*scheduler, prefill(3, 30, 6, 0))); + EXPECT_TRUE(submit(*scheduler, prefill(4, 40, 5, 100))); + + std::vector work = scheduler->get_work(); + EXPECT_EQ(ids(work), (std::vector{1, 2, 3, 4})); + EXPECT_EQ(token_count(work), 1u + 1u + 6u + 5u); + + BatchInput batch = to_batch_input(work); + EXPECT_EQ(batch.size(), 1u + 1u + 6u + 5u); + EXPECT_EQ(batch.inputs.size(), 4u); +} + +// --- decode scheduling ----------------------------------------------------- + +TEST(DecodeTest, ServedInArrivalOrderUpToTheCap) { + SchedulerPtr scheduler = make_scheduler(2, 4); + EXPECT_TRUE(submit(*scheduler, decode(1, 10))); + EXPECT_TRUE(submit(*scheduler, decode(2, 20))); + EXPECT_TRUE(submit(*scheduler, decode(3, 30))); + + EXPECT_EQ(ids(scheduler->get_work()), (std::vector{1, 2})); + EXPECT_TRUE(scheduler->has_work()); + EXPECT_EQ(ids(scheduler->get_work()), (std::vector{3})); + EXPECT_FALSE(scheduler->has_work()); +} + +// A shorter queue must not give a later arrival a head start. +TEST(DecodeTest, StaysFifoAcrossDrainAndRefill) { + SchedulerPtr scheduler = make_scheduler(2, 4); + EXPECT_TRUE(submit(*scheduler, decode(1, 10))); + EXPECT_TRUE(submit(*scheduler, decode(2, 20))); + EXPECT_TRUE(submit(*scheduler, decode(3, 30))); + EXPECT_EQ(ids(scheduler->get_work()), (std::vector{1, 2})); + + EXPECT_TRUE(submit(*scheduler, decode(4, 40))); + EXPECT_EQ(ids(scheduler->get_work()), (std::vector{3, 4})); +} + +TEST(DecodeTest, BeatsPrefillAndStillLeavesRoomForAFullChunk) { + SchedulerPtr scheduler = make_scheduler(3, 4); // batch = 11 + EXPECT_TRUE(submit(*scheduler, decode(1, 10))); + EXPECT_TRUE(submit(*scheduler, decode(2, 20))); + EXPECT_TRUE(submit(*scheduler, decode(3, 30))); + EXPECT_TRUE(submit(*scheduler, prefill(50, 90, 4, 0))); + + std::vector work = scheduler->get_work(); + EXPECT_EQ(ids(work), (std::vector{1, 2, 3, 50})); + EXPECT_EQ(token_count(work), 3u + 4u); +} + +// --- prefill scheduling ---------------------------------------------------- + +TEST(PrefillTest, SubmittedPromptChunksStayInOrder) { + SchedulerPtr scheduler = make_scheduler(1, 2); + std::vector prompt; + prompt.push_back(prefill(1, 7, 2, 0, /*produce_output=*/false)); + prompt.push_back(prefill(2, 7, 2, 2)); + ASSERT_TRUE(scheduler->submit(std::move(prompt))); + + EXPECT_EQ(ids(scheduler->get_work()), (std::vector{1, 2})); +} + +TEST(PrefillTest, LongPromptCannotHogTheBatch) { + SchedulerPtr scheduler = make_scheduler(2, 4); + for (int chunk = 0; chunk < 4; ++chunk) { + EXPECT_TRUE(submit( + *scheduler, + prefill(1 + chunk, 10, 4, static_cast(chunk * 4)))); + } + EXPECT_TRUE(submit(*scheduler, prefill(5, 20, 4, 0))); + EXPECT_TRUE(submit(*scheduler, prefill(6, 20, 4, 4))); + + EXPECT_EQ(ids(scheduler->get_work()), (std::vector{1, 5})); + EXPECT_EQ(ids(scheduler->get_work()), (std::vector{2, 6})); + EXPECT_EQ(ids(scheduler->get_work()), (std::vector{3, 4})) + << "session 20 drained, so session 10 may take two chunks"; + EXPECT_FALSE(scheduler->has_work()); +} + +TEST(PrefillTest, LoneSessionFillsTheBatchAcrossPasses) { + SchedulerPtr scheduler = make_scheduler(1, 4); // batch = 9 + EXPECT_TRUE(submit(*scheduler, prefill(1, 10, 4, 0, false))); + EXPECT_TRUE(submit(*scheduler, prefill(2, 10, 4, 4))); + EXPECT_EQ(ids(scheduler->get_work()), (std::vector{1, 2})); +} + +TEST(PrefillTest, StopsWhenTheNextChunkDoesNotFit) { + SchedulerPtr scheduler = make_scheduler(1, 4); // batch = 9 + EXPECT_TRUE(submit(*scheduler, prefill(1, 10, 4, 0))); + EXPECT_TRUE(submit(*scheduler, prefill(2, 20, 4, 0))); + EXPECT_TRUE(submit(*scheduler, prefill(3, 30, 4, 0))); + + EXPECT_EQ(ids(scheduler->get_work()), (std::vector{1, 2})); + EXPECT_TRUE(scheduler->has_work()); + EXPECT_EQ(ids(scheduler->get_work()), (std::vector{3})); +} + +// A session skipped on size was reached; one the pass never got to was not. +// Both took nothing, so they must keep their original relative order. +TEST(PrefillTest, DeferredSessionOutranksOneNeverReached) { + SchedulerPtr scheduler = make_scheduler(1, 4); // batch = 9 + EXPECT_TRUE(submit(*scheduler, prefill(1, 10, 4, 0))); // 9 -> 5 + EXPECT_TRUE(submit(*scheduler, prefill(2, 20, 3, 0))); // 5 -> 2 + EXPECT_TRUE(submit(*scheduler, prefill(3, 30, 4, 0))); // deferred + EXPECT_TRUE(submit(*scheduler, prefill(4, 40, 2, 0))); // 2 -> 0 + EXPECT_TRUE(submit(*scheduler, prefill(5, 50, 4, 0))); // not reached + + EXPECT_EQ(ids(scheduler->get_work()), (std::vector{1, 2, 4})); + EXPECT_EQ(ids(scheduler->get_work()), (std::vector{3, 5})); +} + +TEST(PrefillTest, DeferredSessionsKeepTheirOrder) { + SchedulerPtr scheduler = make_scheduler(1, 4); // batch = 9 + EXPECT_TRUE(submit(*scheduler, prefill(1, 10, 4, 0))); + EXPECT_TRUE(submit(*scheduler, prefill(2, 20, 4, 0))); + EXPECT_TRUE(submit(*scheduler, prefill(3, 30, 4, 0))); + EXPECT_TRUE(submit(*scheduler, prefill(4, 40, 3, 0))); + + EXPECT_EQ(ids(scheduler->get_work()), (std::vector{1, 2})); + EXPECT_EQ(ids(scheduler->get_work()), (std::vector{3, 4})); +} + +// Two long prompts must not starve a third: served sessions rotate to the back. +TEST(PrefillTest, RotationIsFairAcrossCalls) { + SchedulerPtr scheduler = make_scheduler(1, 4); + TaskId task_id = 1; + for (SessionId session : {10, 20, 30}) { + for (int chunk = 0; chunk < 8; ++chunk) { + EXPECT_TRUE(submit( + *scheduler, + prefill(task_id++, session, 4, static_cast(chunk * 4)))); + } + } + + std::map served; + for (int call = 0; call < 9; ++call) { + for (const Task& task : scheduler->get_work()) { + served[task.input.sid]++; + } + } + EXPECT_EQ(served[10], 6); + EXPECT_EQ(served[20], 6); + EXPECT_EQ(served[30], 6); +} + +TEST(PrefillTest, SessionRejoinsTheRotationAfterDraining) { + SchedulerPtr scheduler = make_scheduler(1, 4); + EXPECT_TRUE(submit(*scheduler, prefill(1, 10, 4, 0))); + EXPECT_EQ(ids(scheduler->get_work()), (std::vector{1})); + EXPECT_TRUE(scheduler->get_work().empty()); + + EXPECT_TRUE(submit(*scheduler, prefill(2, 10, 4, 4))); + EXPECT_EQ(ids(scheduler->get_work()), (std::vector{2})); + + EXPECT_TRUE(submit(*scheduler, prefill(3, 10, 4, 8))); + EXPECT_TRUE(submit(*scheduler, prefill(4, 20, 4, 0))); + EXPECT_EQ(ids(scheduler->get_work()), (std::vector{3, 4})); +} + +// --- what one batch may hold for one session ------------------------------- + +// The executor is promised ranges that adjoin and a single produce_output per +// session. Consecutive prefill chunks satisfy that; two decodes do not. +TEST(BatchCompositionTest, OneDecodePerSessionPerBatch) { + SchedulerPtr scheduler = make_scheduler(8, 4); + EXPECT_TRUE(submit(*scheduler, decode(1, 10, 0))); + EXPECT_TRUE(submit(*scheduler, decode(2, 10, 1))); + EXPECT_TRUE(submit(*scheduler, decode(3, 10, 2))); + EXPECT_TRUE(submit(*scheduler, decode(4, 20, 0))); + + std::vector first = scheduler->get_work(); + EXPECT_EQ(ids(first), (std::vector{1, 4})); + EXPECT_TRUE(at_most_one_task_per_session(first)); + + EXPECT_EQ(ids(scheduler->get_work()), (std::vector{2})) + << "deferred decodes keep their arrival order"; + EXPECT_EQ(ids(scheduler->get_work()), (std::vector{3})); + EXPECT_FALSE(scheduler->has_work()); +} + +TEST(BatchCompositionTest, SecondDecodeForASessionRunsInTheNextBatch) { + SchedulerPtr scheduler = make_scheduler(2, 4); + EXPECT_TRUE(submit(*scheduler, decode(1, 10, 5))); + EXPECT_TRUE(submit(*scheduler, decode(2, 10, 5))); + + EXPECT_EQ(ids(scheduler->get_work()), (std::vector{1})); + EXPECT_EQ(ids(scheduler->get_work()), (std::vector{2})); + EXPECT_FALSE(scheduler->has_work()); +} + +// A decode and a prefill chunk for one session are two ranges that do not +// adjoin, and both ask to produce output. +TEST(BatchCompositionTest, PrefillWaitsWhileTheSessionHasADecode) { + SchedulerPtr scheduler = make_scheduler(8, 4); + EXPECT_TRUE(submit(*scheduler, decode(1, 10, 0))); + EXPECT_TRUE(submit(*scheduler, prefill(2, 10, 4, 8))); + EXPECT_TRUE(submit(*scheduler, prefill(3, 20, 4, 0))); + + std::vector first = scheduler->get_work(); + EXPECT_EQ(ids(first), (std::vector{1, 3})); + EXPECT_TRUE(at_most_one_task_per_session(first)); + + EXPECT_EQ(ids(scheduler->get_work()), (std::vector{2})); +} + +// Several chunks of one prompt are allowed together, because they adjoin and +// only the last asks for output. +TEST(BatchCompositionTest, ConsecutivePrefillChunksMaySharePlaces) { + SchedulerPtr scheduler = make_scheduler(1, 4); // batch = 9 + std::vector prompt; + prompt.push_back(prefill(1, 10, 4, 0, /*produce_output=*/false)); + prompt.push_back(prefill(2, 10, 4, 4)); + ASSERT_TRUE(scheduler->submit(std::move(prompt))); + + std::vector work = scheduler->get_work(); + EXPECT_EQ(ids(work), (std::vector{1, 2})); + EXPECT_TRUE(at_most_one_output_per_session(work)); +} + +// --- cancellation ---------------------------------------------------------- + +TEST(CancelTest, DropsEveryQueuedTaskForTheSessionAndReturnsThem) { + SchedulerPtr scheduler = make_scheduler(2, 4); + std::vector prompt; + prompt.push_back(prefill(1, 77, 4, 0, /*produce_output=*/false)); + prompt.push_back(prefill(2, 77, 4, 4)); + ASSERT_TRUE(scheduler->submit(std::move(prompt))); + + std::vector dropped = scheduler->cancel(77); + EXPECT_EQ(ids(dropped), (std::vector{1, 2})); + for (const Task& task : dropped) { + EXPECT_TRUE(task.cancelled); + } + EXPECT_FALSE(scheduler->has_work()); + EXPECT_TRUE(scheduler->get_work().empty()); +} + +TEST(CancelTest, DropsEveryQueuedDecodeForTheSession) { + SchedulerPtr scheduler = make_scheduler(4, 4); + EXPECT_TRUE(submit(*scheduler, decode(1, 10, 0))); + EXPECT_TRUE(submit(*scheduler, decode(2, 20, 0))); + EXPECT_TRUE(submit(*scheduler, decode(3, 10, 1))); + + EXPECT_EQ(sorted_ids(scheduler->cancel(10)), (std::vector{1, 3})) + << "every queued decode for the session must be dropped, not just one"; + EXPECT_EQ(ids(scheduler->get_work()), (std::vector{2})); + EXPECT_FALSE(scheduler->has_work()); +} + +TEST(CancelTest, LeavesOtherSessionsRunnable) { + SchedulerPtr scheduler = make_scheduler(2, 4); + EXPECT_TRUE(submit(*scheduler, decode(1, 10))); + EXPECT_TRUE(submit(*scheduler, decode(2, 20))); + + EXPECT_EQ(ids(scheduler->cancel(10)), (std::vector{1})); + EXPECT_TRUE(scheduler->has_work()); + EXPECT_EQ(ids(scheduler->get_work()), (std::vector{2})); +} + +TEST(CancelTest, UnknownSessionIsIgnored) { + SchedulerPtr scheduler = make_scheduler(1, 4); + EXPECT_TRUE(scheduler->cancel(404).empty()); + EXPECT_FALSE(scheduler->has_work()); +} + +TEST(CancelTest, DoubleCancelReportsEachTaskOnce) { + SchedulerPtr scheduler = make_scheduler(4, 4); + EXPECT_TRUE(submit(*scheduler, decode(1, 10, 0))); + EXPECT_TRUE(submit(*scheduler, prefill(2, 10, 4, 8))); + + EXPECT_EQ(sorted_ids(scheduler->cancel(10)), (std::vector{1, 2})); + EXPECT_TRUE(scheduler->cancel(10).empty()); + EXPECT_FALSE(scheduler->has_work()); +} + +// A task handed out by get_work() is no longer owned by the scheduler. +TEST(CancelTest, DispatchedTaskIsNoLongerTracked) { + SchedulerPtr scheduler = make_scheduler(1, 4); + EXPECT_TRUE(submit(*scheduler, prefill(9, 77, 4, 12))); + ASSERT_EQ(ids(scheduler->get_work()), (std::vector{9})); + + EXPECT_TRUE(scheduler->cancel(77).empty()); + EXPECT_FALSE(scheduler->has_work()); +} + +TEST(CancelTest, CancellingDispatchedTaskDoesNotHideQueuedWork) { + SchedulerPtr scheduler = make_scheduler(1, 4); + EXPECT_TRUE(submit(*scheduler, decode(1, 10))); + ASSERT_EQ(ids(scheduler->get_work()), (std::vector{1})); + EXPECT_TRUE(submit(*scheduler, decode(2, 20))); + + EXPECT_TRUE(scheduler->cancel(10).empty()); + EXPECT_TRUE(scheduler->has_work()); + EXPECT_EQ(ids(scheduler->get_work()), (std::vector{2})); +} + +// A cancelled entry stays in its deque until scheduling reaches it. Skipping it +// must not disturb live tasks before or after it. +TEST(CancelTest, CancelledDecodeInTheMiddleOfTheQueue) { + SchedulerPtr scheduler = make_scheduler(1, 4); + EXPECT_TRUE(submit(*scheduler, decode(1, 10))); + EXPECT_TRUE(submit(*scheduler, decode(2, 20))); + EXPECT_TRUE(submit(*scheduler, decode(3, 30))); + + EXPECT_EQ(ids(scheduler->cancel(20)), (std::vector{2})); + EXPECT_EQ(ids(scheduler->get_work()), (std::vector{1})); + EXPECT_EQ(ids(scheduler->get_work()), (std::vector{3})); + EXPECT_FALSE(scheduler->has_work()); +} + +TEST(CancelTest, CancelledPrefillSessionLeavesTheRotation) { + SchedulerPtr scheduler = make_scheduler(1, 4); + std::vector first_prompt; + first_prompt.push_back(prefill(1, 10, 4, 0, /*produce_output=*/false)); + first_prompt.push_back(prefill(2, 10, 4, 4)); + ASSERT_TRUE(scheduler->submit(std::move(first_prompt))); + EXPECT_TRUE(submit(*scheduler, prefill(3, 20, 4, 0))); + + EXPECT_EQ(ids(scheduler->cancel(10)), (std::vector{1, 2})); + EXPECT_EQ(ids(scheduler->get_work()), (std::vector{3})); +} + +// The batch in between retires the cancelled session's rotation slot, so +// rejoining after one is the easy case. +TEST(CancelTest, SessionRejoinsAfterAnInterveningBatch) { + SchedulerPtr scheduler = make_scheduler(1, 4); + EXPECT_TRUE(submit(*scheduler, prefill(1, 10, 4, 0))); + ASSERT_EQ(ids(scheduler->cancel(10)), (std::vector{1})); + EXPECT_TRUE(scheduler->get_work().empty()); + + EXPECT_TRUE(submit(*scheduler, prefill(2, 10, 4, 0))); + EXPECT_TRUE(submit(*scheduler, prefill(3, 20, 4, 0))); + EXPECT_EQ(ids(scheduler->get_work()), (std::vector{2, 3})); +} + +// Regression: cancelling used to erase the session's queue while leaving its +// slot in the rotation, so rejoining took a second slot and the session got two +// turns per pass, starving everyone else. There must be no get_work() between +// the cancel and the resubmit, since that would retire the stale slot and hide +// the bug. +TEST(CancelTest, RejoiningAfterCancelDoesNotTakeTwoTurnsPerPass) { + SchedulerPtr scheduler = make_scheduler(1, 4); // batch = 9 + EXPECT_TRUE(submit(*scheduler, prefill(1, 10, 4, 0))); + ASSERT_EQ(ids(scheduler->cancel(10)), (std::vector{1})); + + EXPECT_TRUE(submit(*scheduler, prefill(2, 10, 4, 0, false))); + EXPECT_TRUE(submit(*scheduler, prefill(3, 10, 4, 4))); + EXPECT_TRUE(submit(*scheduler, prefill(4, 20, 4, 0, false))); + EXPECT_TRUE(submit(*scheduler, prefill(5, 20, 4, 4))); + + std::vector work = scheduler->get_work(); + EXPECT_EQ(sessions(work), (std::vector{10, 20})) + << "one chunk each, not two for the session that rejoined"; + EXPECT_EQ(ids(work), (std::vector{2, 4})); +} + +// --- clear ----------------------------------------------------------------- + +TEST(ClearTest, ReturnsAndRemovesAllQueuedTasks) { + SchedulerPtr scheduler = make_scheduler(1, 4); + EXPECT_TRUE(submit(*scheduler, decode(1, 10))); + EXPECT_TRUE(submit(*scheduler, prefill(2, 20, 4, 0))); + + std::vector dropped = scheduler->clear(); + EXPECT_EQ(sorted_ids(dropped), (std::vector{1, 2})); + for (const Task& task : dropped) { + EXPECT_TRUE(task.cancelled); + } + EXPECT_FALSE(scheduler->has_work()); + EXPECT_TRUE(scheduler->get_work().empty()); + EXPECT_TRUE(scheduler->clear().empty()); +} + +TEST(ClearTest, DoesNotClaimTasksAlreadyHandedToTheCaller) { + SchedulerPtr scheduler = make_scheduler(2, 4); + EXPECT_TRUE(submit(*scheduler, decode(1, 10))); + EXPECT_TRUE(submit(*scheduler, decode(2, 20))); + ASSERT_EQ(ids(scheduler->get_work()), (std::vector{1, 2})); + + EXPECT_TRUE(submit(*scheduler, decode(3, 30))); + EXPECT_TRUE(submit(*scheduler, prefill(4, 40, 4, 0))); + EXPECT_EQ(sorted_ids(scheduler->clear()), (std::vector{3, 4})); + EXPECT_TRUE(scheduler->cancel(10).empty()); + EXPECT_TRUE(scheduler->cancel(20).empty()); + EXPECT_FALSE(scheduler->has_work()); +} + +// --- has_work -------------------------------------------------------------- + +TEST(HasWorkTest, EmptySchedulerHasNothingToDo) { + SchedulerPtr scheduler = make_scheduler(1, 4); + EXPECT_FALSE(scheduler->has_work()); + EXPECT_TRUE(scheduler->get_work().empty()); +} + +TEST(HasWorkTest, TracksQueuedTasksOnly) { + SchedulerPtr scheduler = make_scheduler(1, 4); + EXPECT_FALSE(scheduler->has_work()); + + EXPECT_TRUE(submit(*scheduler, decode(1, 10))); + EXPECT_TRUE(scheduler->has_work()); + + scheduler->get_work(); + EXPECT_FALSE(scheduler->has_work()) << "handed-out work is caller-owned"; + EXPECT_TRUE(scheduler->cancel(10).empty()); + EXPECT_FALSE(scheduler->has_work()); +} + +// --- randomized invariants ------------------------------------------------- + +// Every accepted task must be returned exactly once, either from get_work() or +// from cancel(). The public has_work() state must agree with that model, and +// no batch may owe one session two outputs. +TEST(InvariantTest, AcceptedTasksAreNeverLostOrReturnedTwice) { + std::mt19937 rng(1234); + SchedulerPtr scheduler = make_scheduler(3, 4); + + std::map waiting; + TaskId next_task = 1; + SessionId next_session = 1; + int submitted = 0; + int dispatched = 0; + int cancelled = 0; + + for (int operation = 0; operation < 4000; ++operation) { + ASSERT_EQ(scheduler->has_work(), !waiting.empty()) + << "has_work() diverged from the model at operation " << operation; + + switch (rng() % 3) { + case 0: { // Submit one decode or a group of prompt chunks. + const SessionId session = next_session++; + const bool is_decode = rng() % 3 == 0; + const int count = is_decode ? 1 : static_cast(1 + rng() % 3); + std::vector tasks; + for (int i = 0; i < count; ++i) { + const TaskId task_id = next_task++; + const auto position = static_cast(i * 4); + // Only a prompt's last chunk asks for output, as the runner builds + // them, so a batch may legitimately hold several of one session's + // chunks while still owing it a single output. + tasks.push_back( + is_decode ? decode(task_id, session, position) + : prefill( + task_id, + session, + 2 + rng() % 3, + position, + /*produce_output=*/i == count - 1)); + waiting.emplace(task_id, session); + submitted++; + } + ASSERT_TRUE(scheduler->submit(std::move(tasks))); + break; + } + case 1: { // Dispatch a batch. + std::vector work = scheduler->get_work(); + EXPECT_TRUE(at_most_one_output_per_session(work)) + << "a batch owed one session two outputs at operation " + << operation; + for (const Task& task : work) { + EXPECT_FALSE(task.cancelled); + EXPECT_EQ(waiting.erase(task.tid), 1u) + << "a batch returned a task not present in the model"; + dispatched++; + } + break; + } + case 2: { // Cancel a session with work still waiting. + if (waiting.empty()) { + break; + } + auto selected = waiting.begin(); + std::advance( + selected, static_cast(rng() % waiting.size())); + const SessionId session = selected->second; + std::vector expected; + for (const auto& entry : waiting) { + if (entry.second == session) { + expected.push_back(entry.first); + } + } + + std::vector dropped = scheduler->cancel(session); + EXPECT_EQ(sorted_ids(dropped), expected); + for (const Task& task : dropped) { + EXPECT_TRUE(task.cancelled); + EXPECT_EQ(waiting.erase(task.tid), 1u); + cancelled++; + } + break; + } + default: + break; + } + } + + while (scheduler->has_work()) { + std::vector work = scheduler->get_work(); + ASSERT_FALSE(work.empty()) << "has_work() but get_work() made no progress"; + for (const Task& task : work) { + EXPECT_EQ(waiting.erase(task.tid), 1u); + dispatched++; + } + } + + EXPECT_TRUE(waiting.empty()); + EXPECT_FALSE(scheduler->has_work()); + EXPECT_EQ(dispatched + cancelled, submitted); +} + +// --- concurrency ----------------------------------------------------------- + +TEST(ConcurrencyTest, ProducersAndConsumerMakeProgressWithoutLoss) { + SchedulerPtr scheduler = make_scheduler(4, 8); + constexpr int kProducers = 4; + constexpr int kPerProducer = 250; + std::atomic stop{false}; + std::atomic next_task{1}; + std::atomic accepted{0}; + std::atomic dispatched{0}; + + std::thread consumer([&] { + while (!stop.load() || scheduler->has_work()) { + std::vector work = scheduler->get_work(); + dispatched += static_cast(work.size()); + if (work.empty()) { + std::this_thread::yield(); + } + } + }); + + std::vector producers; + for (int producer = 0; producer < kProducers; ++producer) { + producers.emplace_back([&] { + for (int i = 0; i < kPerProducer; ++i) { + const TaskId task_id = next_task.fetch_add(1); + const auto position = static_cast(i); + Task task = (i % 3 == 0) ? decode(task_id, task_id, position) + : prefill(task_id, task_id, 4, position); + if (submit(*scheduler, std::move(task))) { + accepted++; + } + } + }); + } + for (std::thread& producer : producers) { + producer.join(); + } + + const auto deadline = + std::chrono::steady_clock::now() + std::chrono::seconds(5); + while (dispatched.load() < accepted.load() && + std::chrono::steady_clock::now() < deadline) { + std::this_thread::yield(); + } + stop.store(true); + consumer.join(); + + EXPECT_EQ(accepted.load(), kProducers * kPerProducer) + << "a valid, uniquely identified task was rejected"; + EXPECT_EQ(dispatched.load(), accepted.load()) + << "an accepted task was never dispatched"; + EXPECT_FALSE(scheduler->has_work()); +} + +TEST(ConcurrencyTest, ObserversAreSafeDuringScheduling) { + SchedulerPtr scheduler = make_scheduler(2, 4); + std::atomic stop{false}; + std::atomic observing{false}; + std::atomic observations{0}; + + std::thread observer([&] { + observing.store(true); + while (!stop.load()) { + (void)scheduler->has_work(); + (void)scheduler->max_batch_tokens(); + observations++; + } + }); + while (!observing.load()) { + std::this_thread::yield(); + } + + for (TaskId task_id = 1; task_id <= 500; ++task_id) { + ASSERT_TRUE(submit(*scheduler, decode(task_id, task_id))); + (void)scheduler->get_work(); + } + stop.store(true); + observer.join(); + + EXPECT_GT(observations.load(), 0) << "observer never ran"; + EXPECT_FALSE(scheduler->has_work()); +} diff --git a/extension/llm/batching/types.h b/extension/llm/batching/types.h new file mode 100644 index 00000000000..46bf5e58ee0 --- /dev/null +++ b/extension/llm/batching/types.h @@ -0,0 +1,103 @@ +/* + * 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 + +// The vocabulary shared by the runner, the scheduler, and the executor. +// +// An Input is one slice of work for one session, either a decode token or one +// chunk of a prompt, never a whole generation. A Task is an Input plus the +// scheduling identity used to order and cancel it. + +#include +#include +#include +#include + +namespace executorch { +namespace extension { +namespace llm { +namespace batching { + +using Token = std::int64_t; +using SessionId = std::int64_t; +using Position = std::int32_t; +// Wide enough that a monotonically issued id cannot wrap in any realistic +// lifetime, so ids never have to be recycled. +using TaskId = std::int64_t; + +// Sampling policy for a generation. Installed on the session before its tasks +// are submitted, so it does not ride on every Input. +struct SamplingParams { + float temperature = 0.0f; + float top_p = 1.0f; + std::int32_t top_k = 0; +}; + +struct Input { + SessionId sid; + bool produce_output; + + // The selected slice is tokens[offset : offset + size]. It starts at the + // absolute logical position `position + offset`; `position` is the base of + // the complete backing vector, not of the slice. + size_t offset; + size_t size; + + std::shared_ptr> tokens; + Position position; +}; + +struct Output { + SessionId sid; + + // What this input produced: usually one token, can be more than one for + // speculative decoding + std::vector tokens; +}; + +struct Task { + TaskId tid; + bool cancelled; + Input input; + bool is_decode; +}; + +struct BatchInput { + std::vector inputs; + size_t size() const { + size_t sz = 0; + for (const auto& i : inputs) { + sz += i.size; + } + return sz; + } +}; + +// The executor's view of a batch: the Inputs, without the tid, cancelled flag, +// and is_decode that only the runner and scheduler use. +// +// Moves each Input out of its Task, preserving task order, so outputs[i] +// answers batch.inputs[i]. +inline BatchInput to_batch_input(std::vector& tasks) { + BatchInput batch; + batch.inputs.reserve(tasks.size()); + for (Task& t : tasks) { + batch.inputs.push_back(std::move(t.input)); + } + return batch; +} + +struct BatchOutput { + std::vector> outputs; +}; + +} // namespace batching +} // namespace llm +} // namespace extension +} // namespace executorch