Introduce support for dynamic batching

examples/python/engine/continuous-batching.py

@@ -55,7 +55,7 @@ def __init__(
         tokenizer: og.Tokenizer,
         engine: og.Engine,
         num_requests: int,
-        load_factor: float = 0.2,
+        load_factor: float = 1,
         debug: bool = False,
     ):
         self.model = model
@@ -135,6 +135,7 @@ def run(args: argparse.Namespace):
         engine.tokenizer,
         engine.engine,
         args.num_requests,
+        load_factor=args.load_factor,
         debug=args.debug,
     )
 
@@ -146,7 +147,7 @@ def run(args: argparse.Namespace):
     end = time.time()
 
     request_pool.bar.close()
-    print(f"⌛Tokens per second: {engine.tokens_decoded / (end - start):.2f}")
+    print(f"⌛ Tokens per second: {engine.tokens_decoded / (end - start):.2f}")
 
 
 if __name__ == "__main__":
@@ -181,6 +182,13 @@ def run(args: argparse.Namespace):
         default=1,
         help="Number of requests to process in the pool",
     )
+    parser.add_argument(
+        "-l",
+        "--load_factor",
+        type=float,
+        default=1.0,
+        help="Load factor to control the number of preloaded in-flight requests (default: 1.0)",
+    )
     args = parser.parse_args()
 
     run(args)

src/config.cpp

@@ -313,6 +313,12 @@ struct DecoderInputs_Element : JSON::Element {
       v_.past_key_values_length = JSON::Get<std::string_view>(value);
     } else if (name == "cache_indirection") {
       v_.cache_indirection = JSON::Get<std::string_view>(value);
+    } else if (name == "cumulative_sequence_lengths") {
+      v_.cumulative_sequence_lengths = JSON::Get<std::string_view>(value);
+    } else if (name == "past_sequence_lengths") {
+      v_.past_sequence_lengths = JSON::Get<std::string_view>(value);
+    } else if (name == "block_table") {
+      v_.block_table = JSON::Get<std::string_view>(value);
     } else {
       throw JSON::unknown_value_error{};
     }
@@ -883,6 +889,67 @@ struct Search_Element : JSON::Element {
   Config::Search& v_;
 };
 
+struct DynamicBatching_Element : JSON::Element {
+  explicit DynamicBatching_Element(std::optional<Config::Engine::DynamicBatching>& v) : v_{v} {}
+
+  void OnValue(std::string_view name, JSON::Value value) override {
+    if (!v_)
+      v_ = Config::Engine::DynamicBatching{};
+
+    if (name == "block_size") {
+      v_->block_size = static_cast<size_t>(JSON::Get<double>(value));
+    } else if (name == "num_blocks") {
+      v_->num_blocks = static_cast<size_t>(JSON::Get<double>(value));
+    } else if (name == "gpu_utilization_factor") {
+      v_->gpu_utilization_factor = static_cast<float>(JSON::Get<double>(value));
+    } else if (name == "max_batch_size") {
+      v_->max_batch_size = static_cast<size_t>(JSON::Get<double>(value));
+    } else {
+      throw JSON::unknown_value_error{};
+    }
+  }
+
+ private:
+  std::optional<Config::Engine::DynamicBatching>& v_;
+};
+
+struct StaticBatching_Element : JSON::Element {
+  explicit StaticBatching_Element(std::optional<Config::Engine::StaticBatching>& v) : v_{v} {}
+
+  void OnValue(std::string_view name, JSON::Value value) override {
+    if (name == "max_batch_size") {
+      v_->max_batch_size = static_cast<size_t>(JSON::Get<double>(value));
+    } else {
+      throw JSON::unknown_value_error{};
+    }
+  }
+
+ private:
+  std::optional<Config::Engine::StaticBatching>& v_;
+};
+
+struct Engine_Element : JSON::Element {
+  explicit Engine_Element(Config::Engine& v) : v_{v} {}
+
+  Element& OnObject(std::string_view name) override {
+    if (name == "dynamic_batching") {
+      if (v_.static_batching)
+        v_.static_batching.reset();
+      return dynamic_batching_;
+    } else if (name == "static_batching") {
+      if (v_.dynamic_batching)
+        v_.dynamic_batching.reset();
+      return static_batching_;
+    }
+    throw JSON::unknown_value_error{};
+  }
+
+ private:
+  Config::Engine& v_;
+  DynamicBatching_Element dynamic_batching_{v_.dynamic_batching};
+  StaticBatching_Element static_batching_{v_.static_batching};
+};
+
 void SetSearchNumber(Config::Search& search, std::string_view name, double value) {
   try {
     Search_Element(search).OnValue(name, value);
@@ -1045,12 +1112,16 @@ struct Root_Element : JSON::Element {
     if (name == "search") {
       return search_element_;
     }
+    if (name == "engine") {
+      return engine_element_;
+    }
     throw JSON::unknown_value_error{};
   }
 
   Config& config_;
   Model_Element model_element_{config_.model};
   Search_Element search_element_{config_.search};
+  Engine_Element engine_element_{config_.engine};
 };
 
 struct RootObject_Element : JSON::Element {

src/config.h

@@ -275,6 +275,21 @@ struct Config {
     int random_seed{-1};               // -1 = Seed with random device, otherwise use value to seed RNG
   } search;
 
+  struct Engine {
+    struct DynamicBatching {
+      size_t block_size{256};                       // Total number of slots per block.
+      std::optional<size_t> num_blocks;             // Total number of blocks per layer.
+      std::optional<float> gpu_utilization_factor;  // Fraction of free GPU memory to use for key-value cache.
+      size_t max_batch_size{16};                    // Maximum batch size for dynamically batching requests.
+    };
+    std::optional<DynamicBatching> dynamic_batching;  // Dynamic batching settings
+
+    struct StaticBatching {
+      size_t max_batch_size{4};  // Maximum batch size for static batching
+    };
+    std::optional<StaticBatching> static_batching;  // Static batching settings
+  } engine;                                         // Engine settings
+
   void AddMapping(const std::string& nominal_name, const std::string& graph_name);
   // Returns graph name and true if the nominal name is found in the mapping
   // otherwise returns the nominal name and false

src/cuda/interface.cpp

@@ -172,6 +172,10 @@ struct CudaInterfaceImplBase : DeviceInterface {
   void LaunchAddLogitsMask(float* batch_logits, int batch_beam_size, int vocab_size, const uint32_t* logits_mask) override {
     cuda::LaunchAddLogitsMask(batch_logits, batch_beam_size, vocab_size, logits_mask, GetStream());
   }
+
+  void GetAvailableMemory(size_t& free_bytes, size_t& total_bytes) override {
+    cudaMemGetInfo(&free_bytes, &total_bytes);
+  }
 };
 
 struct CudaInterfaceImpl final : CudaInterfaceImplBase {

src/engine/block.cpp

@@ -0,0 +1,102 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "../generators.h"
+#include "block.h"
+
+#include <numeric>
+#include <algorithm>
+
+namespace Generators {
+
+Block::Block(size_t id, size_t slots, size_t block_size)
+    : id_(id), size_(slots), capacity_(block_size) {}
+
+size_t Block::Id() const {
+  return id_;
+}
+
+size_t Block::Size() const {
+  return size_;
+}
+
+bool Block::IsFull() const {
+  return Size() == Capacity();
+}
+
+size_t Block::EmptySlots() const {
+  return Capacity() - Size();
+}
+
+size_t Block::Capacity() const {
+  return capacity_;
+}
+
+void Block::AddSlot() {
+  if (IsFull()) {
+    throw std::runtime_error("Cannot add a slot. The block is full.");
+  }
+
+  size_++;
+}
+
+std::vector<size_t> Block::SlotIds() const {
+  std::vector<size_t> slot_ids(Size(), 0);
+  std::iota(slot_ids.begin(), slot_ids.end(), Id() * Capacity());
+  return slot_ids;
+}
+
+BlockPool::BlockPool(size_t block_size, size_t num_blocks)
+    : block_size_(block_size), capacity_(num_blocks) {}
+
+std::vector<std::shared_ptr<Block>> BlockPool::AllocateBlocks(size_t num_slots) {
+  const auto allocate_block = [this](size_t num_slots) {
+    for (size_t i = 0; i < Capacity(); ++i) {
+      if (blocks_[i] == nullptr) {
+        blocks_[i] = std::make_shared<Block>(i, num_slots, block_size_);
+        return blocks_[i];
+      }
+    }
+    return std::shared_ptr<Block>();
+  };
+
+  if (BlocksNeeded(num_slots) > AvailableBlocks()) {
+    throw std::runtime_error("Requested number of blocks " + std::to_string(BlocksNeeded(num_slots)) +
+                             " for number of slots " + std::to_string(num_slots) +
+                             " exceeds available blocks " + std::to_string(AvailableBlocks()) + ".");
+  }
+
+  std::vector<std::shared_ptr<Block>> allocated_blocks;
+  for (size_t i = 0; i < num_slots; i += block_size_) {
+    auto block = allocate_block(std::min(block_size_, num_slots - i));
+    if (!block) {
+      throw std::runtime_error("Failed to allocate a block.");
+    }
+    allocated_blocks.push_back(block);
+  }
+  return allocated_blocks;
+}
+
+void BlockPool::Free(const std::vector<std::shared_ptr<Block>>& blocks) {
+  for (const auto& block : blocks) {
+    blocks_[block->Id()].reset();
+  }
+}
+
+size_t BlockPool::AvailableBlocks() const {
+  return std::count_if(blocks_.begin(), blocks_.end(), [](const std::shared_ptr<Block>& block) { return block == nullptr; });
+}
+
+size_t BlockPool::Size() const {
+  return Capacity() - AvailableBlocks();
+}
+
+size_t BlockPool::Capacity() const {
+  return capacity_;
+}
+
+size_t BlockPool::BlocksNeeded(size_t num_slots) {
+  return (num_slots + block_size_ - 1) / block_size_;
+}
+
+}  // namespace Generators

src/engine/block.h

@@ -0,0 +1,65 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include <memory>
+#include <vector>
+
+namespace Generators {
+
+/*
+ * Block represents a contiguous set of slots in the paged key-value cache.
+ * Each block has a fixed capacity (number of slots it can hold) and tracks
+ * the number of currently used slots.
+ */
+struct Block {
+  Block(size_t id, size_t slots, size_t block_size);
+
+  size_t Id() const;
+
+  size_t Size() const;
+
+  bool IsFull() const;
+
+  size_t Capacity() const;
+
+  size_t EmptySlots() const;
+
+  void AddSlot();
+
+  std::vector<size_t> SlotIds() const;
+
+ private:
+  size_t id_;
+  size_t size_;
+  size_t capacity_;
+};
+
+/*
+ * BlockPool manages a pool of blocks for the paged key-value cache.
+ * It allows allocation and deallocation of blocks, and keeps track
+ * of the total capacity and currently available blocks.
+ */
+struct BlockPool {
+  BlockPool(size_t block_size, size_t num_blocks);
+
+  size_t AvailableBlocks() const;
+
+  size_t Size() const;
+
+  size_t Capacity() const;
+
+  std::vector<std::shared_ptr<Block>> AllocateBlocks(size_t num_slots);
+
+  void Free(const std::vector<std::shared_ptr<Block>>& blocks);
+
+  size_t BlocksNeeded(size_t num_slots);
+
+ private:
+  const size_t block_size_;
+  const size_t capacity_;
+  std::vector<std::shared_ptr<Block>> blocks_{capacity_};
+};
+
+}  // namespace Generators