Implementing Dynamic Batching for LLM Inference in Custom C++ Server Gateways — code snippets

snippet-1.txt

#include <queue>
#include <mutex>
#include <condition_variable>
#include <chrono>
#include <vector>
#include <optional>

template <typename T>
class ConcurrentBatchQueue {
public:
    ConcurrentBatchQueue() = default;
    ~ConcurrentBatchQueue() = default;

    // Disable copy and assignment to prevent accidental overhead
    ConcurrentBatchQueue(const ConcurrentBatchQueue&) = delete;
    ConcurrentBatchQueue& operator=(const ConcurrentBatchQueue&) = delete;

    void Push(T&& item) {
        {
            std::lock_guard<std::mutex> lock(mutex_);
            queue_.push(std::move(item));
        }
        cv_.notify_one();
    }

    // Block-waits until the queue contains items or the timeout expires.
    // Once active, it pops up to max_items in a single lock acquisition.
    std::vector<T> PopBatch(size_t max_items, std::chrono::microseconds timeout) {
        std::unique_lock<std::mutex> lock(mutex_);
        
        if (queue_.empty()) {
            cv_.wait_for(lock, timeout, [this] { return !queue_.empty(); });
        }
        
        std::vector<T> batch;
        if (queue_.empty()) {
            return batch; // Timeout expired, returned empty batch
        }

        size_t items_to_pop = std::min(queue_.size(), max_items);
        batch.reserve(items_to_pop);
        
        while (!queue_.empty() && batch.size() < max_items) {
            batch.push_back(std::move(queue_.front()));
            queue_.pop();
        }
        return batch;
    }

    size_t Size() const {
        std::lock_guard<std::mutex> lock(mutex_);
        return queue_.size();
    }

    bool Empty() const {
        std::lock_guard<std::mutex> lock(mutex_);
        return queue_.empty();
    }

private:
    std::queue<T> queue_;
    mutable std::mutex mutex_;
    std::condition_variable cv_;
};

snippet-2.txt

#include <vector>
#include <string>
#include <functional>
#include <future>
#include <chrono>
#include <memory>

struct InferenceRequest {
    std::string request_id;
    std::vector<int32_t> input_token_ids;
    uint32_t max_new_tokens = 512;
    float temperature = 0.7f;
    float top_p = 0.9f;
    
    // Callback invoked asynchronously as each token is generated
    std::function<void(const std::string& token_text, bool is_finished)> response_callback;
    
    // Shared promise to notify the ingress gRPC worker of completion or failure
    std::shared_ptr<std::promise<void>> completion_promise;
    
    // SLA tracking metrics
    std::chrono::steady_clock::time_point arrival_time;
};

snippet-3.txt

#include <thread>
#include <atomic>
#include <iostream>

class DynamicBatcher {
public:
    DynamicBatcher(ConcurrentBatchQueue<InferenceRequest>& queue, 
                   size_t max_batch_size, 
                   size_t max_tokens, 
                   std::chrono::microseconds max_delay)
        : queue_(queue), max_batch_size_(max_batch_size), 
          max_tokens_in_batch_(max_tokens), max_queue_delay_(max_delay), 
          running_(false) {}

    ~DynamicBatcher() { Stop(); }

    void Start() {
        if (!running_.exchange(true)) {
            batcher_thread_ = std::thread(&DynamicBatcher::Loop, this);
        }
    }

    void Stop() {
        if (running_.exchange(false)) {
            if (batcher_thread_.joinable()) {
                batcher_thread_.join();
            }
        }
    }

private:
    void Loop() {
        while (running_) {
            // Retrieve a batch. Wait up to max_queue_delay_ if queue is empty.
            auto batch = queue_.PopBatch(max_batch_size_, max_queue_delay_);
            if (batch.empty()) {
                continue; // Yield and recheck run state
            }

            std::vector<InferenceRequest> active_batch;
            size_t current_tokens = 0;

            for (auto& req : batch) {
                size_t req_tokens = req.input_token_ids.size();
                // Check if adding this request violates batch token capacity
                if (current_tokens + req_tokens > max_tokens_in_batch_) {
                    if (active_batch.empty()) {
                        // Edge case: single request exceeds limits. Force run it alone.
                        active_batch.push_back(std::move(req));
                        break;
                    } else {
                        // Requeue the request to preserve FIFO ordering
                        queue_.Push(std::move(req));
                    }
                } else {
                    current_tokens += req_tokens;
                    active_batch.push_back(std::move(req));
                }
            }

            if (!active_batch.empty()) {
                DispatchToEngine(std::move(active_batch));
            }
        }
    }

    void DispatchToEngine(std::vector<InferenceRequest> batch) {
        // Direct call to low-overhead C++ engine bindings
        // e.g., ModelEngine::ExecuteBatchAsync()
        std::cout << "[Batcher] Dispatched " << batch.size() 
                  << " requests containing " << GetTotalTokens(batch) 
                  << " total tokens to GPU engine.\n";
    }

    size_t GetTotalTokens(const std::vector<InferenceRequest>& batch) const {
        size_t total = 0;
        for (const auto& r : batch) total += r.input_token_ids.size();
        return total;
    }

    ConcurrentBatchQueue<InferenceRequest>& queue_;
    size_t max_batch_size_;
    size_t max_tokens_in_batch_;
    std::chrono::microseconds max_queue_delay_;
    std::atomic<bool> running_;
    std::thread batcher_thread_;
};

snippet-4.txt

#include <vector>
#include <memory>
#include <algorithm>
#include <stdexcept>

enum class SequenceState {
    PREFILL,
    DECODE,
    FINISHED
};

// Mock interfaces for engine dependencies
class KVCacheManager;
class ModelEngine {
public:
    struct ModelBatch {};
    virtual std::vector<int32_t> ExecuteStep(const ModelBatch& batch) = 0;
    virtual ~ModelEngine() = default;
};
class Tokenizer {
public:
    virtual std::string Decode(int32_t token_id) = 0;
    virtual ~Tokenizer() = default;
};

struct ActiveSequence {
    InferenceRequest request;
    SequenceState state = SequenceState::PREFILL;
    size_t generated_tokens = 0;
    std::vector<int32_t> output_tokens;
    std::vector<uint32_t> kv_block_ids;
};

class ContinuousScheduler {
public:
    ContinuousScheduler(ConcurrentBatchQueue<InferenceRequest>& queue,
                        std::unique_ptr<KVCacheManager> kv_mgr,
                        std::unique_ptr<ModelEngine> engine,
                        std::unique_ptr<Tokenizer> tokenizer,
                        size_t max_batch_size)
        : queue_(queue), kv_cache_manager_(std::move(kv_mgr)),
          model_engine_(std::move(engine)), tokenizer_(std::move(tokenizer)),
          max_batch_size_(max_batch_size) {}

    void RunSchedulerIteration() {
        // Step 1: Evict sequences marked finished and resolve their futures
        active_sequences_.erase(
            std::remove_if(active_sequences_.begin(), active_sequences_.end(),
                [this](const auto& seq) {
                    if (seq->state == SequenceState::FINISHED) {
                        seq->request.completion_promise->set_value();
                        FreeKVCacheForSequence(seq);
                        return true;
                    }
                    return false;
                }),
            active_sequences_.end()
        );

        // Step 2: Ingest new requests up to available batch slots
        size_t available_slots = max_batch_size_ - active_sequences_.size();
        while (available_slots > 0 && !queue_.Empty()) {
            auto popped_batch = queue_.PopBatch(1, std::chrono::microseconds(0));
            if (popped_batch.empty()) {
                break;
            }
            
            auto seq = std::make_shared<ActiveSequence>(ActiveSequence{std::move(popped_batch[0])});
            size_t prefill_tokens = seq->request.input_token_ids.size();

            // Consult KV Cache block allocator before admitting
            if (AllocateKVCacheForPrefill(seq, prefill_tokens)) {
                active_sequences_.push_back(seq);
                available_slots--;
            } else {
                // Requeue due to VRAM resource exhaustion
                queue_.Push(std::move(seq->request));
                break; // Stop scheduling new prefils to prevent cache thrashing
            }
        }

        if (active_sequences_.empty()) {
            return;
        }

        // Step 3: Run exactly one token generation step
        ModelEngine::ModelBatch batch_payload = BuildBatchPayload();
        std::vector<int32_t> next_tokens = model_engine_->ExecuteStep(batch_payload);

        // Step 4: Update state of all active sequences
        for (size_t i = 0; i < active_sequences_.size(); ++i) {
            auto& seq = active_sequences_[i];
            int32_t token = next_tokens[i];
            
            seq->output_tokens.push_back(token);
            seq->generated_tokens++;

            bool is_eos = (token == eos_token_id_);
            bool limit_reached = (seq->generated_tokens >= seq->request.max_new_tokens);

            std::string token_str = tokenizer_->Decode(token);
            seq->request.response_callback(token_str, is_eos || limit_reached);

            if (is_eos || limit_reached) {
                seq->state = SequenceState::FINISHED;
            } else {
                seq->state = SequenceState::DECODE;
                // Allocate single token slot in KV Cache block table for the next decode iteration
                if (!AllocateKVCacheForDecode(seq)) {
                    // Cache exhausted mid-generation: trigger pre-emption / fallback logic
                    seq->state = SequenceState::FINISHED;
                }
            }
        }
    }

private:
    bool AllocateKVCacheForPrefill(std::shared_ptr<ActiveSequence>& seq, size_t tokens);
    bool AllocateKVCacheForDecode(std::shared_ptr<ActiveSequence>& seq);
    void FreeKVCacheForSequence(const std::shared_ptr<ActiveSequence>& seq);
    ModelEngine::ModelBatch BuildBatchPayload();

    ConcurrentBatchQueue<InferenceRequest>& queue_;
    std::vector<std::shared_ptr<ActiveSequence>> active_sequences_;
    std::unique_ptr<KVCacheManager> kv_cache_manager_;
    std::unique_ptr<ModelEngine> model_engine_;
    std::unique_ptr<Tokenizer> tokenizer_;
    size_t max_batch_size_;
    int32_t eos_token_id_ = 2; // Default for models like Llama
};

snippet-5.txt

#include <vector>
#include <cstdint>
#include <stdexcept>
#include <mutex>
#include <atomic>

class KVCacheManager {
public:
    KVCacheManager(uint32_t total_physical_blocks, uint32_t block_size_tokens)
        : block_size_(block_size_tokens), total_blocks_(total_physical_blocks), 
          free_blocks_(total_physical_blocks) {
        
        free_block_list_.reserve(total_physical_blocks);
        for (uint32_t i = 0; i < total_physical_blocks; ++i) {
            free_block_list_.push_back(i);
        }
    }

    ~KVCacheManager() = default;

    // Evaluate if there are enough free blocks to host the sequence tokens
    bool CanAllocate(size_t num_tokens) const {
        std::lock_guard<std::mutex> lock(mutex_);
        uint32_t blocks_needed = ComputeBlocksNeeded(num_tokens);
        return free_block_list_.size() >= blocks_needed;
    }

    // Allocate blocks and return their physical indexes
    std::vector<uint32_t> Allocate(size_t num_tokens) {
        std::lock_guard<std::mutex> lock(mutex_);
        uint32_t blocks_needed = ComputeBlocksNeeded(num_tokens);

        if (free_block_list_.size() < blocks_needed) {
            throw std::runtime_error("[KV Cache] Allocation failed: out of physical block memory.");
        }

        std::vector<uint32_t> allocated_indices;
        allocated_indices.reserve(blocks_needed);
        
        for (uint32_t i = 0; i < blocks_needed; ++i) {
            allocated_indices.push_back(free_block_list_.back());
            free_block_list_.pop_back();
        }
        
        free_blocks_ -= blocks_needed;
        return allocated_indices;
    }

    // Return physical blocks to the free pool
    void Free(const std::vector<uint32_t>& blocks) {
        std::lock_guard<std::mutex> lock(mutex_);
        for (uint32_t block_id : blocks) {
            free_block_list_.push_back(block_id);
        }
        free_blocks_ += blocks.size();
    }

    size_t FreeBlockCount() const {
        return free_blocks_.load();
    }

    size_t TotalBlockCount() const {
        return total_blocks_;
    }

private:
    uint32_t ComputeBlocksNeeded(size_t num_tokens) const {
        return (num_tokens + block_size_ - 1) / block_size_;
    }

    const uint32_t block_size_;
    const uint32_t total_blocks_;
    std::atomic<uint32_t> free_blocks_;
    std::vector<uint32_t> free_block_list_;
    mutable std::mutex mutex_;
};