Repository: ai-dynamo/dynamo
Pull request #1033: feat: vllm mock workers, Rusty skeleton
Author: login: PeaBrane
Labels: name: size/XXL
name: feat
Created: 2025-05-11T18:07:34Z | Merged: 2025-05-21T09:46:28Z

### Overview:

Implements a mock worker in Rust simulating a vllm-ish behavior. The core components for now are:

1. A synchronous KV manager implementing a simple LRU eviction strategy, largely following [this implementation](https://docs.vllm.ai/en/stable/design/v1/prefix_caching.html)
2. An asynchronous scheduler + simulator all in one that uses the said KV manager, along with other budget considerations like `max_num_batched_tokens`
3. The prefill compute is assumed to scale roughly with `new_tokens * (new_tokens + cached_tokens)` scaled by a dummy magic number

To limit the scope of this PR, it is not hooked up to a mock `AsyncEngine` or dynamo endpoint yet. Neither are any Python bindings written. But currently, a mock worker can be launched and some meaningful `FowardPassMetrics` can be generated (as generated by actual vllm workers for KV routing)

### Where to start reviewing

The core logics are in:

1. `KvManager.process()` in `mocker/kv_manager.rs`, containing the logic for handling the 4 `MoveBlock` variants (see below)
2. The background event loop in `Scheduler.new()` in `mocker/scheduler.rs`, handling receiving a request, scheduling a request, and simulating the generation process

### Motivation

is two-fold:
1. To perform predictive routing, based on more accurate predictions of the `KvIndexer`, beyond the [current heuristics](https://github.com/ai-dynamo/dynamo/blob/main/lib/llm/src/kv_router/scheduler.rs#L220).
3. To implement mock vllm workers for at-scale testings. Most likely use case will be to launch multiple mock vllm workers at once.

### Implementation Details

#### Move Blocks

There is a `MoveBlock` enum with three variants that can be sent around as events, all handled synchronously by the KV manager

1. `Use`: First checks if block is in active pool; if so, increment reference count. Next checks if in inactive pool, if so, move to active. Lastly, try evicting from the inactive pool to make room. If inactive pool is empty, then pre-empt the oldest running request.
2. `Destroy`: Simply removes the block from the active pool. This is used as `vllm` does not cache partial blocks.
4. `Deref`: Decrements the reference count of a block in active pool by one. If reference count is zero, move block into inactive pool.
5. `Promote`: Promotes a partial block (identified with uuid) into a full block (identified with a global block hash)

Note `Use` adds blocks from root to leaf, and `Destroy` and `Deref` removes blocks from leaf to root.

#### Evictor

is a modification of the lazy heap introduced in [this vllm PR ](https://github.com/vllm-project/vllm/pull/7209). The gist is as follows:

1. Use a `VecDeque` / queue to maintain the blocks, with the order guaranteed by the user when pushing (old blocks first, and leaves first)
2. Do not update the queue on timestamp update (e.g. when a block is referenced / touched), since it is O(n), just leave the stale entry there. Instead, a separate hash map is used to keep the updated entries.
6. During removal / eviction, the stale entries will be naturally evicted, as they are "older" than the current ones
7. If the heap grows beyond a certain threshold (defaults to 50), rebuild the `VecDeque` from the hashmap

May be memory intensive if `evict` is rarely called (should probably use a `BTreeSet`)

### Limitation

1. Does not support chunked prefill, but should not be difficult to add.
2. Only consider KV budget and batched token budget during scheduling (does not consider say `max_requests`, which rarely is the bottleneck)

### Integration

Will make a near-future effort to integrate with existing components `tokens.rs` and the recent `block_manager`. May also make sense to not use too many existing components as a stand-alone mock API. Open to discussion.


Diff:
diff --git a/lib/llm/src/lib.rs b/lib/llm/src/lib.rs
index 2bf3d804d082..122a1d6e66af 100644
--- a/lib/llm/src/lib.rs
+++ b/lib/llm/src/lib.rs
@@ -28,6 +28,7 @@ pub mod hub;
 pub mod key_value_store;
 pub mod kv_router;
 pub use kv_router::DEFAULT_KV_BLOCK_SIZE;
+pub mod mocker;
 pub mod model_card;
 pub mod model_type;
 pub mod preprocessor;
diff --git a/lib/llm/src/mocker.rs b/lib/llm/src/mocker.rs
new file mode 100644
index 000000000000..2a9e63a9e2fa
--- /dev/null
+++ b/lib/llm/src/mocker.rs
@@ -0,0 +1,20 @@
+// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+pub mod evictor;
+pub mod kv_manager;
+pub mod protocols;
+pub mod scheduler;
+pub mod sequence;
diff --git a/lib/llm/src/mocker/evictor.rs b/lib/llm/src/mocker/evictor.rs
new file mode 100644
index 000000000000..47a312eede4f
--- /dev/null
+++ b/lib/llm/src/mocker/evictor.rs
@@ -0,0 +1,191 @@
+// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::cmp::Eq;
+use std::collections::{HashMap, VecDeque};
+use std::hash::Hash;
+use std::time::Instant;
+
+/// An LRU evictor that maintains objects and evicts them based on their
+/// last accessed time. Implements a "lazy" eviction mechanism where:
+/// 1. The priority queue does not immediately reflect updates or removes
+/// 2. Objects are pushed to the queue in order of increasing priority (older objects first)
+/// 3. The user must ensure objects are added in correct priority (temporal order)
+/// 4. Remove and update operations are lazy - entries remain in the queue until
+///    they are either evicted or cleaned up during maintenance
+#[derive(Debug)]
+pub struct LRUEvictor<T: Clone + Eq + Hash> {
+    free_table: HashMap<T, f64>,
+    priority_queue: VecDeque<(T, f64)>,
+    cleanup_threshold: usize,
+    start_time: Instant,
+}
+
+impl<T: Clone + Eq + Hash> Default for LRUEvictor<T> {
+    fn default() -> Self {
+        Self {
+            free_table: HashMap::new(),
+            priority_queue: VecDeque::new(),
+            cleanup_threshold: 50,
+            start_time: Instant::now(),
+        }
+    }
+}
+
+impl<T: Clone + Eq + Hash> LRUEvictor<T> {
+    /// Create a new LRUEvictor with the default cleanup threshold
+    pub fn new(cleanup_threshold: usize) -> Self {
+        Self {
+            cleanup_threshold,
+            ..Default::default()
+        }
+    }
+
+    /// Get the current timestamp as seconds since initialization
+    pub fn current_timestamp(&self) -> f64 {
+        self.start_time.elapsed().as_secs_f64()
+    }
+
+    /// Get an iterator over the keys in the evictor
+    pub fn keys(&self) -> std::collections::hash_map::Keys<'_, T, f64> {
+        self.free_table.keys()
+    }
+
+    /// Insert or update an object in the evictor with current timestamp
+    pub fn insert(&mut self, object: T) {
+        let timestamp = self.current_timestamp();
+        self._insert(object, timestamp);
+    }
+
+    /// Check if the evictor contains the given object
+    pub fn contains(&self, object: &T) -> bool {
+        self.free_table.contains_key(object)
+    }
+
+    /// Evict an object based on LRU policy
+    /// Returns the evicted object or None if no objects are available
+    pub fn evict(&mut self) -> Option<T> {
+        if self.free_table.is_empty() {
+            return None;
+        }
+
+        while let Some((object, last_accessed)) = self.priority_queue.pop_front() {
+            let Some(&current_last_accessed) = self.free_table.get(&object) else {
+                continue; // entry is already removed
+            };
+
+            if current_last_accessed == last_accessed {
+                self.free_table.remove(&object);
+                return Some(object);
+            } // otherwise entry is stale
+        }
+
+        None
+    }
+
+    /// Insert or update an object in the evictor
+    fn _insert(&mut self, object: T, last_accessed: f64) {
+        self.free_table.insert(object.clone(), last_accessed);
+        self.priority_queue.push_back((object, last_accessed));
+        self.cleanup_if_necessary();
+    }
+
+    /// Remove an object from the evictor
+    /// We don't remove from the priority queue immediately, as that would be inefficient
+    /// Outdated entries will be filtered out during eviction or cleanup
+    pub fn remove(&mut self, object: &T) -> bool {
+        self.free_table.remove(object).is_some()
+    }
+
+    /// Get the number of objects in the evictor
+    pub fn len(&self) -> usize {
+        self.free_table.len()
+    }
+
+    /// Check if the evictor is empty
+    pub fn is_empty(&self) -> bool {
+        self.free_table.is_empty()
+    }
+
+    /// Check if cleanup is necessary and perform it if needed
+    fn cleanup_if_necessary(&mut self) {
+        if self.priority_queue.len() > self.cleanup_threshold * self.free_table.len() {
+            self.cleanup();
+        }
+    }
+
+    /// Clean up the priority queue by removing outdated entries
+    fn cleanup(&mut self) {
+        let mut new_priority_queue = VecDeque::new();
+        for (object, timestamp) in self.priority_queue.drain(..) {
+            let Some(&current_timestamp) = self.free_table.get(&object) else {
+                continue;
+            };
+
+            if current_timestamp == timestamp {
+                new_priority_queue.push_back((object, timestamp));
+            }
+        }
+        self.priority_queue = new_priority_queue;
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use rstest::rstest;
+
+    #[rstest]
+    #[case(1)]
+    #[case(2)]
+    #[case(3)]
+    fn test_lru_evictor_eviction_order(#[case] threshold: usize) {
+        // Create a new LRUEvictor with the given cleanup threshold
+        let mut evictor = LRUEvictor::<i32>::new(threshold);
+
+        // Add items in the specified order with small delays between each
+        evictor.insert(4);
+        std::thread::sleep(std::time::Duration::from_millis(1));
+        evictor.insert(3);
+        std::thread::sleep(std::time::Duration::from_millis(1));
+        evictor.insert(2);
+        std::thread::sleep(std::time::Duration::from_millis(1));
+        evictor.insert(1);
+        std::thread::sleep(std::time::Duration::from_millis(1));
+        evictor.insert(5);
+        std::thread::sleep(std::time::Duration::from_millis(1));
+        evictor.insert(1); // Updates timestamp for 1
+        std::thread::sleep(std::time::Duration::from_millis(1));
+        evictor.insert(4); // Updates timestamp for 4
+        std::thread::sleep(std::time::Duration::from_millis(1));
+        evictor.insert(2); // Updates timestamp for 2
+
+        // Verify the eviction order
+        println!("Testing with threshold {}", threshold);
+        let evicted = evictor.evict().unwrap();
+        assert_eq!(evicted, 3);
+        let evicted = evictor.evict().unwrap();
+        assert_eq!(evicted, 5);
+        let evicted = evictor.evict().unwrap();
+        assert_eq!(evicted, 1);
+        let evicted = evictor.evict().unwrap();
+        assert_eq!(evicted, 4);
+        let evicted = evictor.evict().unwrap();
+        assert_eq!(evicted, 2);
+        let evicted = evictor.evict();
+        assert_eq!(evicted, None);
+        assert_eq!(evictor.len(), 0);
+    }
+}
diff --git a/lib/llm/src/mocker/kv_manager.rs b/lib/llm/src/mocker/kv_manager.rs
new file mode 100644
index 000000000000..8a7a8fefed0f
--- /dev/null
+++ b/lib/llm/src/mocker/kv_manager.rs
@@ -0,0 +1,427 @@
+// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! # KV Manager
+//! A synchronous implementation of a block manager that handles MoveBlock signals for caching KV blocks.
+//!
+//! ## Block Operations
+//! The KV manager processes four types of MoveBlock signals:
+//!
+//! ### Use
+//! - Checks if block exists in active pool → increment reference count
+//! - If in inactive pool → move to active pool
+//! - If neither → try evicting from inactive pool to make room
+//! - If inactive pool is empty → pre-empt the oldest running request
+//!
+//! ### Destroy
+//! - Removes the block from the active pool
+//!
+//! ### Deref
+//! - Decrements reference count of a block in active pool
+//! - If count reaches zero → move block to inactive pool
+//!
+//! ### Promote
+//! - Converts a partial block (uuid) into a full block (global block hash)
+//!
+//! ## Preemption
+//! If a Use operation fails (typically due to insufficient space), a false boolean signal
+//! is returned to the scheduler for preemption. Initial KV block allocations for new requests
+//! should not fail due to the watermark checking.
+//!
+//! ## NOTE
+//! For simplicity (or non-simplicity), reference counting is tracked manually instead of using
+//! the more idiomatic built-in Arc reference counter. This can be considered a shadow / mirror
+//! implementation of the main block manager.
+
+use crate::mocker::evictor::LRUEvictor;
+use crate::mocker::protocols::{MoveBlock, PrefillCost, UniqueBlock};
+use crate::mocker::sequence::ActiveSequence;
+use derive_getters::Getters;
+use std::collections::{HashMap, HashSet};
+
+#[derive(Getters)]
+pub struct KvManager {
+    #[getter(copy)]
+    max_capacity: usize,
+
+    #[getter(copy)]
+    block_size: usize,
+
+    active_blocks: HashMap<UniqueBlock, usize>,
+
+    inactive_blocks: LRUEvictor<UniqueBlock>,
+
+    all_blocks: HashSet<UniqueBlock>,
+}
+
+impl KvManager {
+    pub fn new(max_capacity: usize, block_size: usize) -> Self {
+        let active_blocks = HashMap::new();
+        let inactive_blocks = LRUEvictor::default();
+        let all_blocks = HashSet::new();
+
+        KvManager {
+            max_capacity,
+            block_size,
+            active_blocks,
+            inactive_blocks,
+            all_blocks,
+        }
+    }
+
+    /// Process a MoveBlock instruction synchronously
+    pub fn process(&mut self, event: &MoveBlock) -> bool {
+        match event {
+            MoveBlock::Use(hashes, _) => {
+                for hash in hashes {
+                    // First check if it already exists in active blocks
+                    if let Some(ref_count) = self.active_blocks.get_mut(hash) {
+                        // Block already active, just increment reference count
+                        *ref_count += 1;
+                        continue;
+                    }
+
+                    // Then check if it exists in inactive and move it to active if found
+                    if self.inactive_blocks.remove(hash) {
+                        // Insert into active with reference count 1
+                        self.active_blocks.insert(hash.clone(), 1);
+                        continue;
+                    }
+
+                    // Get counts for capacity check
+                    let active_count = self.active_blocks.len();
+                    let inactive_count = self.inactive_blocks.len();
+
+                    // If at max capacity, evict the oldest entry from inactive blocks
+                    if active_count + inactive_count >= self.max_capacity {
+                        if let Some(evicted) = self.inactive_blocks.evict() {
+                            // Remove evicted block from all_blocks
+                            self.all_blocks.remove(&evicted);
+                        } else {
+                            // Cannot evict block, meaning no free blocks left in inactive pool
+                            // Send a signal, scheduler would expect to handle preemption upon receiving this
+                            return false;
+                        }
+                    }
+
+                    // Now insert the new block in active blocks with reference count 1
+                    self.active_blocks.insert(hash.clone(), 1);
+                    // Add to all_blocks as it's a new block
+                    self.all_blocks.insert(hash.clone());
+                }
+            }
+            MoveBlock::Destroy(hashes) => {
+                // Loop in inverse direction
+                for hash in hashes.iter().rev() {
+                    self.active_blocks.remove(hash).unwrap();
+                    // Remove from all_blocks when destroyed
+                    assert!(self.all_blocks.remove(hash));
+                }
+            }
+            MoveBlock::Deref(hashes) => {
+                // Loop in inverse direction
+                for hash in hashes.iter().rev() {
+                    // Decrement reference count and check if we need to move to inactive
+                    if let Some(ref_count) = self.active_blocks.get_mut(hash) {
+                        if *ref_count == 0 {
+                            panic!("Negative reference count would be encountered after Deref.");
+                        }
+                        *ref_count -= 1;
+
+                        // If reference count reaches zero, remove from active and move to inactive
+                        if *ref_count == 0 {
+                            self.active_blocks.remove(hash);
+                            // Use the LRUEvictor's timing functionality
+                            self.inactive_blocks.insert(hash.clone());
+                        }
+                    }
+                }
+            }
+            MoveBlock::Promote(uuid, hash) => {
+                let uuid_block = UniqueBlock::PartialBlock(*uuid);
+                let hash_block = UniqueBlock::FullBlock(*hash);
+
+                let Some(ref_count) = self.active_blocks.remove(&uuid_block) else {
+                    let in_all_blocks = self.all_blocks.contains(&uuid_block);
+                    panic!(
+                        "Missing active block for promotion: {:?}. Block still exists: {}",
+                        uuid_block, in_all_blocks
+                    );
+                };
+
+                // Replace with hash block, keeping the same reference count
+                self.active_blocks.insert(hash_block.clone(), ref_count);
+
+                // Update all_blocks
+                assert!(self.all_blocks.remove(&uuid_block));
+                self.all_blocks.insert(hash_block);
+            }
+        }
+
+        // Return true if we made it this far
+        true
+    }
+
+    /// Get the count of blocks in the input list that aren't in all_blocks
+    pub fn probe_new_blocks(&self, blocks: &[UniqueBlock]) -> usize {
+        blocks
+            .iter()
+            .filter(|&block| !self.all_blocks.contains(block))
+            .count()
+    }
+
+    /// Get the current capacity (active blocks + inactive blocks)
+    pub fn current_capacity(&self) -> usize {
+        let active = self.active_blocks.len();
+        let inactive = self.inactive_blocks.len();
+        active + inactive
+    }
+
+    /// Get the current capacity as a percentage of the maximum capacity
+    pub fn current_capacity_perc(&self) -> f64 {
+        let current = self.current_capacity() as f64;
+        current / self.max_capacity as f64
+    }
+
+    /// Get the number of active blocks
+    pub fn num_active_blocks(&self) -> usize {
+        self.active_blocks.len()
+    }
+
+    /// Get the number of inactive blocks
+    pub fn num_inactive_blocks(&self) -> usize {
+        self.inactive_blocks.len()
+    }
+
+    /// Get the keys of inactive blocks
+    pub fn get_inactive_blocks(&self) -> Vec<&UniqueBlock> {
+        self.inactive_blocks.keys().collect()
+    }
+
+    /// Get the keys of active blocks
+    pub fn get_active_blocks(&self) -> Vec<&UniqueBlock> {
+        self.active_blocks.keys().collect()
+    }
+
+    /// Check if a sequence can be scheduled and calculate cost if possible
+    pub fn try_schedule(
+        &self,
+        sequence: &ActiveSequence,
+        watermark: f64,
+        tokens_budget: usize,
+    ) -> Option<PrefillCost> {
+        // Return None immediately if tokens_budget is 0
+        if tokens_budget == 0 {
+            return None;
+        }
+
+        // Get unique blocks from the sequence
+        let unique_blocks = sequence.unique_blocks();
+
+        // Get the count of new blocks
+        let new_blocks = self.probe_new_blocks(unique_blocks);
+
+        // Calculate current usage and available capacity
+        let active_count = self.active_blocks.len();
+
+        // Check if we can schedule based on the watermark
+        if (active_count + new_blocks) as f64 > (1.0 - watermark) * self.max_capacity as f64 {
+            return None;
+        }
+
+        // Calculate overlap blocks
+        let overlap_blocks = unique_blocks.len() - new_blocks;
+
+        // Calculate new tokens
+        let new_tokens = sequence.num_input_tokens() - overlap_blocks * self.block_size;
+
+        // // Print the full equation with actual values substituted
+        // println!("{} = {} - ({} * {}) (new_tokens = num_input_tokens - overlap_blocks * block_size)",
+        //     new_tokens,
+        //     sequence.num_input_tokens(),
+        //     overlap_blocks,
+        //     self.block_size);
+
+        // Return None if new_tokens exceeds tokens_budget
+        if new_tokens > tokens_budget {
+            return None;
+        }
+
+        // Calculate prefill compute
+        let prefill_compute =
+            new_tokens as f64 * (new_tokens + overlap_blocks * self.block_size) as f64;
+
+        Some(PrefillCost {
+            new_tokens,
+            prefill_compute,
+        })
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_failure_on_max_capacity() {
+        // Create a KvManager with 10 blocks capacity
+        let mut manager = KvManager::new(10, 16);
+
+        // Helper function to use multiple blocks that returns the response
+        fn use_blocks(manager: &mut KvManager, ids: Vec<u64>) -> bool {
+            let blocks = ids.into_iter().map(UniqueBlock::FullBlock).collect();
+            manager.process(&MoveBlock::Use(blocks, None))
+        }
+
+        // First use 10 blocks (0 to 9) in a batch
+        let response = use_blocks(&mut manager, (0..10).collect());
+        assert!(response, "Expected success response");
+
+        // Verify we are at capacity
+        assert_eq!(manager.current_capacity(), 10);
+
+        // The 11th block should return false, not panic
+        let response = use_blocks(&mut manager, vec![10]);
+        assert!(
+            !response,
+            "Expected failure response when exceeding max capacity"
+        );
+    }
+
+    #[test]
+    // This is taken directly from the example in the vllm v1 prefix caching docs
+    fn test_block_lifecycle_stringent() {
+        // Create a KvManager with 10 blocks capacity
+        let mut manager = KvManager::new(10, 16);
+
+        // Helper function to use multiple blocks
+        fn use_blocks(manager: &mut KvManager, ids: Vec<u64>) {
+            let blocks = ids.into_iter().map(UniqueBlock::FullBlock).collect();
+            manager.process(&MoveBlock::Use(blocks, None));
+        }
+
+        // Helper function to destroy multiple blocks
+        fn destroy_blocks(manager: &mut KvManager, ids: Vec<u64>) {
+            let blocks = ids.into_iter().map(UniqueBlock::FullBlock).collect();
+            manager.process(&MoveBlock::Destroy(blocks));
+        }
+
+        // Helper function to deref multiple blocks
+        fn deref_blocks(manager: &mut KvManager, ids: Vec<u64>) {
+            let blocks = ids.into_iter().map(UniqueBlock::FullBlock).collect();
+            manager.process(&MoveBlock::Deref(blocks));
+        }
+
+        // Helper function to check if active blocks contain expected blocks with expected ref counts
+        fn assert_active_blocks(manager: &KvManager, expected_blocks: &[(u64, usize)]) {
+            assert_eq!(
+                manager.active_blocks().len(),
+                expected_blocks.len(),
+                "Active blocks count doesn't match expected"
+            );
+
+            for &(id, ref_count) in expected_blocks {
+                let block = UniqueBlock::FullBlock(id);
+                assert!(
+                    manager.active_blocks().contains_key(&block),
+                    "Block {} not found in active blocks",
+                    id
+                );
+                assert_eq!(
+                    manager.active_blocks().get(&block),
+                    Some(&ref_count),
+                    "Block {} has wrong reference count",
+                    id
+                );
+            }
+        }
+
+        // Helper function to check if inactive blocks contain expected blocks
+        fn assert_inactive_blocks(
+            manager: &KvManager,
+            expected_size: usize,
+            expected_blocks: &[u64],
+        ) {
+            let inactive_blocks = manager.get_inactive_blocks();
+            let inactive_blocks_count = manager.inactive_blocks().len();
+
+            assert_eq!(
+                inactive_blocks_count, expected_size,
+                "Inactive blocks count doesn't match expected"
+            );
+
+            for &id in expected_blocks {
+                let block = UniqueBlock::FullBlock(id);
+                assert!(
+                    inactive_blocks.iter().any(|&b| *b == block),
+                    "Block {} not found in inactive blocks",
+                    id
+                );
+            }
+        }
+
+        // First use blocks 0, 1, 2, 3, 4 in a batch
+        use_blocks(&mut manager, (0..5).collect());
+
+        // Then use blocks 0, 1, 5, 6 in a batch
+        use_blocks(&mut manager, vec![0, 1, 5, 6]);
+
+        // Check that the blocks 0 and 1 are in active blocks, both with reference counts of 2
+        assert_active_blocks(
+            &manager,
+            &[(0, 2), (1, 2), (2, 1), (3, 1), (4, 1), (5, 1), (6, 1)],
+        );
+
+        // Now destroy block 4
+        destroy_blocks(&mut manager, vec![4]);
+
+        // And deref blocks 3, 2, 1, 0 in this order as a batch
+        deref_blocks(&mut manager, vec![0, 1, 2, 3]);
+
+        // Check that the inactive_blocks is size 2 (via num_objects) and contains 3 and 2
+        assert_inactive_blocks(&manager, 2, &[3, 2]);
+        assert_active_blocks(&manager, &[(0, 1), (1, 1), (5, 1), (6, 1)]);
+
+        // Now destroy block 6
+        destroy_blocks(&mut manager, vec![6]);
+
+        // And deref blocks 5, 1, 0 as a batch
+        deref_blocks(&mut manager, vec![0, 1, 5]);
+
+        // Check that the inactive_blocks is size 5, and contains 0, 1, 2, 3, 5
+        assert_inactive_blocks(&manager, 5, &[0, 1, 2, 3, 5]);
+        assert_active_blocks(&manager, &[]);
+
+        // Now use 0, 1, 2, 7, 8, 9 as a batch
+        use_blocks(&mut manager, vec![0, 1, 2, 7, 8, 9]);
+
+        // Check that the inactive_blocks is size 2, and contains 3 and 5
+        assert_inactive_blocks(&manager, 2, &[3, 5]);
+        assert_active_blocks(&manager, &[(0, 1), (1, 1), (2, 1), (7, 1), (8, 1), (9, 1)]);
+
+        // Test the new_blocks method - only block 4 should be new out of [0,1,2,3,4]
+        let blocks_to_check: Vec<UniqueBlock> = vec![0, 1, 2, 3, 4]
+            .into_iter()
+            .map(UniqueBlock::FullBlock)
+            .collect();
+        assert_eq!(manager.probe_new_blocks(&blocks_to_check), 1);
+
+        // Now use blocks 10, 11, 12 as a batch
+        use_blocks(&mut manager, vec![10, 11, 12]);
+
+        // Check that the inactive_blocks is size 1 and contains only 5
+        assert_inactive_blocks(&manager, 1, &[5]);
+    }
+}
diff --git a/lib/llm/src/mocker/protocols.rs b/lib/llm/src/mocker/protocols.rs
new file mode 100644
index 000000000000..2b551db61b9c
--- /dev/null
+++ b/lib/llm/src/mocker/protocols.rs
@@ -0,0 +1,93 @@
+// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use serde::{Deserialize, Serialize};
+use uuid::Uuid;
+
+pub type Token = u32;
+pub type LocalBlockHash = u64;
+/// A global hash identifier for blocks
+pub type GlobalHash = u64;
+pub type NumBlocks = usize;
+
+/// Represents an active block in the cache with a reference count
+#[derive(Debug, Clone, Hash, Eq, PartialEq, Serialize, Deserialize)]
+pub enum UniqueBlock {
+    /// Block identified by UUID
+    PartialBlock(Uuid),
+    /// Block identified by hash
+    FullBlock(GlobalHash),
+}
+
+impl Default for UniqueBlock {
+    fn default() -> Self {
+        // Generate a random UUID when default is used
+        Self::PartialBlock(Uuid::new_v4())
+    }
+}
+
+/// Represents different block movement operations in the cache
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+pub enum MoveBlock {
+    Use(Vec<UniqueBlock>, Option<f64>),
+    Destroy(Vec<UniqueBlock>),
+    Deref(Vec<UniqueBlock>),
+    Promote(Uuid, GlobalHash),
+}
+
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct DirectRequest {
+    pub tokens: Vec<Token>,
+    pub max_output_tokens: usize,
+    pub uuid: Option<Uuid>,
+}
+
+/// Represents the cost of prefilling content in the cache
+#[derive(Debug, Clone, Serialize, Deserialize)]
+pub struct PrefillCost {
+    pub new_tokens: usize,
+    pub prefill_compute: f64,
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_unique_block_default_uniqueness() {
+        // Create 10 default UniqueBlock instances
+        let blocks: Vec<UniqueBlock> = (0..10).map(|_| UniqueBlock::default()).collect();
+
+        // Extract UUIDs from each block
+        let mut uuids = Vec::new();
+        for block in blocks {
+            match block {
+                UniqueBlock::PartialBlock(uuid) => uuids.push(uuid),
+                _ => panic!("Expected UuidIdentifier variant"),
+            }
+        }
+
+        // Check that all UUIDs are unique by comparing each with every other
+        for i in 0..uuids.len() {
+            for j in i + 1..uuids.len() {
+                assert_ne!(
+                    uuids[i], uuids[j],
+                    "UUID at index {} and {} are identical",
+                    i, j
+                );
+            }
+        }
+    }
+}
diff --git a/lib/llm/src/mocker/scheduler.rs b/lib/llm/src/mocker/scheduler.rs
new file mode 100644
index 000000000000..e71647feabb8
--- /dev/null
+++ b/lib/llm/src/mocker/scheduler.rs
@@ -0,0 +1,585 @@
+// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Asynchronous Scheduler for LLM Request Management
+//!
+//! This module implements an asynchronous scheduler that handles three main functions:
+//! 1. Receiving new requests and placing them in the waiting queue
+//! 2. Scheduling waiting requests against available KV cache resources
+//! 3. Simulating the execution of running requests with realistic timing
+//!
+//! ## Scheduling Process
+//! The scheduler uses a watermark-based approach to determine if there's sufficient
+//! KV cache space for new requests. It also enforces a batched tokens budget to prevent
+//! oversubscription of computational resources. Only requests that can be allocated
+//! these resources are moved from waiting to running state.
+//!
+//! ## Request Simulation
+//! The simulation models two key phases:
+//! - Prefill phase: Uses a quadratic cost function: (cached_tokens + new_tokens) * new_tokens
+//! - Decode phase: Uses a cost function proportional to active KV blocks (linear)
+//!
+//! ## Resource Management
+//! The scheduler communicates with the KvManager through MoveBlock signals at each
+//! stage of request processing. When resources become constrained, it employs an
+//! LRU-based preemption strategy where the oldest running request is evicted and
+//! placed at the back of the waiting queue to be rescheduled later.
+//!
+//! ## NOTE
+//! The current prefill and decoding time simulations are not scientific at all and are WIP
+
+use crate::kv_router::protocols::ForwardPassMetrics;
+use crate::mocker::evictor::LRUEvictor;
+use crate::mocker::kv_manager::KvManager;
+use crate::mocker::protocols::DirectRequest;
+use crate::mocker::protocols::{MoveBlock, PrefillCost, UniqueBlock};
+use crate::mocker::sequence::ActiveSequence;
+use std::collections::HashMap;
+use std::collections::VecDeque;
+use std::sync::Arc;
+use tokio::sync::{mpsc, Mutex};
+use tokio::time::{interval, Duration};
+use tokio_util::sync::CancellationToken;
+use uuid::Uuid;
+
+/// Enum representing either a direct request or an active sequence
+pub enum Request {
+    Direct(DirectRequest),
+    Active(ActiveSequence),
+}
+
+#[derive(Default)]
+struct SchedulerState {
+    waiting: VecDeque<Uuid>,
+    ready: VecDeque<Uuid>,
+    running: LRUEvictor<Uuid>,
+    requests: HashMap<Uuid, Request>,
+    prefill_costs: HashMap<Uuid, Option<PrefillCost>>,
+}
+
+impl SchedulerState {
+    /// Create a new UUID for a DirectRequest, add it to requests, and push the UUID to waiting.
+    fn receive(&mut self, request: DirectRequest) -> Uuid {
+        // Use the provided UUID if available, otherwise generate a new one
+        let uuid = request.uuid.unwrap_or_else(Uuid::new_v4);
+
+        // Add the request to the map and waiting queue
+        self.requests.insert(uuid, Request::Direct(request));
+        self.waiting.push_back(uuid);
+        uuid
+    }
+
+    /// Get the next UUID from ready or waiting queue and its associated Request.
+    /// Returns from ready if not empty, otherwise from waiting, or None if both are empty.
+    /// Also removes the Request from the requests HashMap.
+    fn next(&mut self) -> Option<(Uuid, Request)> {
+        let uuid = self
+            .ready
+            .pop_front()
+            .or_else(|| self.waiting.pop_front())?;
+        let request = self.requests.remove(&uuid)?;
+        Some((uuid, request))
+    }
+
+    /// Move a UUID and its Request to the ready queue.
+    fn make_ready(&mut self, uuid: Uuid, active_seq: ActiveSequence) {
+        self.requests.insert(uuid, Request::Active(active_seq));
+        self.ready.push_back(uuid);
+    }
+
+    /// Schedule the request with the given UUID.
+    /// Returns the creation signal from the ActiveSequence.
+    fn run(&mut self, uuid: Uuid, active_seq: ActiveSequence) -> MoveBlock {
+        // Insert the request into the map
+        self.requests.insert(uuid, Request::Active(active_seq));
+
+        // Get the creation signal
+        let Some(Request::Active(sequence)) = self.requests.get(&uuid) else {
+            panic!("Failed to get ActiveSequence for UUID");
+        };
+        let Some(signal) = sequence.creation_signal() else {
+            panic!("Failed to get creation signal from ActiveSequence");
+        };
+
+        // Add to running requests
+        self.running.insert(uuid);
+        signal.clone()
+    }
+
+    /// Set the prefill cost for a UUID
+    fn set_prefill_cost(&mut self, uuid: Uuid, cost: Option<PrefillCost>) {
+        self.prefill_costs.insert(uuid, cost);
+    }
+
+    /// Get the prefill compute value for a UUID if available
+    fn get_prefill_compute(&self, uuid: &Uuid) -> Option<f64> {
+        self.prefill_costs
+            .get(uuid)
+            .and_then(|cost| cost.as_ref())
+            .map(|cost| cost.prefill_compute)
+    }
+
+    /// Calculate the current running batched tokens
+    fn num_batched_tokens(&self) -> usize {
+        self.prefill_costs
+            .values()
+            .map(|cost| match cost {
+                Some(cost) => cost.new_tokens,
+                None => 1,
+            })
+            .sum()
+    }
+
+    /// Remove a UUID and its associated Request from collections.
+    fn complete(&mut self, uuid: &Uuid) {
+        // println!("Request {} will complete", uuid);
+        self.running.remove(uuid);
+        self.requests.remove(uuid);
+        self.prefill_costs.remove(uuid);
+    }
+
+    /// Preempt the oldest running request by evicting it from running, resetting the sequence,
+    /// and adding it back to the waiting queue.
+    /// Returns the signal from reset_with_signal or None if no requests are running.
+    fn preempt(&mut self) -> Option<Vec<MoveBlock>> {
+        // Evict the oldest UUID from running
+        let uuid = self.running.evict()?;
+        eprintln!("Request {} will be preempted", uuid);
+
+        // Remove the request from the requests HashMap and ensure it's an ActiveSequence
+        let request = self.requests.remove(&uuid)?;
+
+        // Remove the prefill cost to force recomputation
+        self.prefill_costs.remove(&uuid);
+
+        // Extract the ActiveSequence from the Request enum
+        let Request::Active(mut active_sequence) = request else {
+            panic!("Expected ActiveSequence in running queue")
+        };
+
+        // Reset the sequence and get the new sequence and signal
+        let signals = active_sequence.reset_with_signal();
+
+        // Insert the new sequence back into the requests map and add to waiting queue
+        self.requests.insert(uuid, Request::Active(active_sequence));
+        self.waiting.push_back(uuid);
+
+        Some(signals)
+    }
+}
+
+/// Manages scheduling of requests using KvManager resources
+#[derive(Clone)]
+pub struct Scheduler {
+    state: Arc<Mutex<SchedulerState>>,
+    kv_manager: Arc<Mutex<KvManager>>,
+    request_tx: mpsc::Sender<DirectRequest>,
+}
+
+impl Scheduler {
+    /// Create a new Scheduler with the given parameters
+    pub fn new(
+        kv_capacity: usize,
+        watermark: f64,
+        block_size: usize,
+        chunk_size: Option<usize>,
+        output_tx: Option<mpsc::Sender<Uuid>>,
+        cancellation_token: Option<CancellationToken>,
+    ) -> Self {
+        // Create KvManager internally
+        let kv_manager = KvManager::new(kv_capacity, block_size);
+
+        let token_capacity: usize = 8192;
+        let state = Arc::new(Mutex::new(SchedulerState::default()));
+
+        let kv_manager = Arc::new(Mutex::new(kv_manager));
+        let chunk_size = chunk_size.unwrap_or(256);
+
+        // Create channel for request handling
+        let (request_tx, mut request_rx) = mpsc::channel::<DirectRequest>(1024);
+
+        // Use provided cancellation token or create new one
+        let cancellation_token = cancellation_token.unwrap_or_default();
+        let token_clone = cancellation_token.clone();
+
+        // Create a clone for the background task
+        let state_clone = state.clone();
+        let kv_manager_clone = kv_manager.clone();
+        let output_tx_clone = output_tx.clone();
+
+        // Spawn main background task with cancellation token
+        tokio::spawn(async move {
+            let mut schedule_interval = interval(Duration::from_millis(5));
+            let mut simulate_interval = interval(Duration::from_millis(1));
+
+            loop {
+                tokio::select! {
+                    biased;
+
+                    // Enqueue new request
+                    Some(request) = request_rx.recv() => {
+                        let mut state = state_clone.lock().await;
+                        state.receive(request);
+                    }
+
+                    // Try Scheduling Requests
+                    _ = schedule_interval.tick() => {
+                        let mut state_guard = state_clone.lock().await;
+                        let mut kv_manager_guard = kv_manager_clone.lock().await;
+
+                        // Process DirectRequests, converting them to ActiveSequence and scheduling them until we can't
+                        // schedule anymore.
+                        while let Some((uuid, request)) = state_guard.next() {
+                            let active_sequence = get_active_sequence(request, block_size, chunk_size);
+
+                            // Calculate token budget using new_tokens from PrefillCost
+                            let total_prefill_tokens = state_guard.num_batched_tokens();
+                            let tokens_budget = token_capacity.saturating_sub(total_prefill_tokens);
+
+                            // Check if it can be scheduled
+                            let Some(prefill_cost) = kv_manager_guard.try_schedule(&active_sequence, watermark, tokens_budget) else {
+                                state_guard.make_ready(uuid, active_sequence);
+                                break;
+                            };
+
+                            // Get creation signal and schedule the request
+                            let signal = state_guard.run(uuid, active_sequence);
+                            kv_manager_guard.process(&signal);
+                            state_guard.set_prefill_cost(uuid, Some(prefill_cost));
+                        }
+                    }
+
+                    // Check for cancellation
+                    _ = token_clone.cancelled() => {
+                        break;
+                    }
+
+                    // Simulate running requests (prefill + decode)
+                    _ = simulate_interval.tick() => {
+                        let mut state_guard = state_clone.lock().await;
+                        let mut kv_manager_guard = kv_manager_clone.lock().await;
+
+                        // Base time needed for decoding (assumed memory bound on KV cache)
+                        let active_tokens = kv_manager_guard.num_active_blocks() * block_size;
+                        // TODO: 2 is a dummy / magic scaling factor
+                        let mut generation_time = Duration::from_micros((active_tokens / 2) as u64);
+
+                        // Process each running request
+                        let uuids: Vec<Uuid> = state_guard.running.keys().cloned().collect();
+                        for uuid in uuids {
+                            // Check if UUID is still in running_requests, if not skip this iteration
+                            if !state_guard.running.contains(&uuid) {
+                                continue;
+                            }
+
+                            // Get prefill compute value first
+                            let prefill_compute = state_guard.get_prefill_compute(&uuid);
+
+                            // Get the active sequence for this UUID
+                            let sequence = state_guard.requests.get_mut(&uuid)
+                                .and_then(|req| if let Request::Active(seq) = req { Some(seq) } else { None })
+                                .expect("UUID in running_requests must have a corresponding active sequence");
+
+                            // Generate token and get signals
+                            let signals = sequence.generate();
+
+                            // Accumulate sleep duration based on prefill_compute if available
+                            // prefill compute = (cached_tokens + new_tokens) * new_tokens
+                            let sleep_ms = if let Some(compute) = prefill_compute {
+                                // TODO: 1024 is a dummy / magic scaling factor
+                                (compute / 1024.0) as u64
+                            } else { 0 };
+                            generation_time += Duration::from_micros(sleep_ms);
+
+                            // Process all signals with the KvManager
+                            // Handling of preemption on failure
+                            if !process_signals(&mut kv_manager_guard, &signals) {
+                                sequence.pop();  // revert the failed generation op
+
+                                // free_signal derefs the preempted blocks
+                                let Some(free_signal) = state_guard.preempt() else {
+                                    panic!("Failed to acquire signal to free KV blocks from preemption");
+                                };
+
+                                for signal in free_signal {
+                                    kv_manager_guard.process(&signal);
+                                }
+                                continue;
+                            }
+
+                            // Send UUID notification for each generated token
+                            // TODO: hook this up to an AsyncEngine
+                            if let Some(tx) = &output_tx_clone {
+                                let _ = tx.try_send(uuid);
+                            }
+
+                            // Check if we're done after generating
+                            if sequence.generated_tokens() >= sequence.max_output_tokens() {
+                                state_guard.complete(&uuid);
+                                continue;
+                            }
+
+                            // Transition to decode (no prefill cost)
+                            if sequence.generated_tokens() == 1 {
+                                state_guard.set_prefill_cost(uuid, None);
+                            }
+                        }
+
+                        // Sleep once for the accumulated duration
+                        if generation_time.as_millis() > 0 {
+                            tokio::time::sleep(generation_time).await;
+                        }
+                    }
+                }
+            }
+        });
+
+        Self {
+            state,
+            kv_manager,
+            request_tx,
+        }
+    }
+
+    /// Add a new request to the waiting queue
+    pub async fn receive(&self, request: DirectRequest) {
+        let _ = self.request_tx.send(request).await;
+    }
+
+    /// Get the count of waiting requests
+    pub async fn waiting_count(&self) -> usize {
+        let state = self.state.lock().await;
+        state.waiting.len()
+    }
+
+    /// Get the count of running requests
+    pub async fn running_count(&self) -> usize {
+        let state = self.state.lock().await;
+        state.running.len()
+    }
+
+    /// Get the current capacity of the KvManager
+    pub async fn kv_usage_perc(&self) -> f64 {
+        let kv_manager = self.kv_manager.lock().await;
+        kv_manager.current_capacity_perc()
+    }
+
+    /// Returns forward pass metrics for monitoring purposes
+    pub async fn get_forward_pass_metrics(&self) -> ForwardPassMetrics {
+        let state = self.state.lock().await;
+        let kv_manager = self.kv_manager.lock().await;
+
+        // Get the active blocks and total capacity from KvManager
+        let active_blocks_count = kv_manager.active_blocks().len() as u64;
+        let total_capacity = kv_manager.max_capacity() as u64;
+
+        // Calculate GPU cache usage percentage
+        let gpu_cache_usage_perc = if total_capacity > 0 {
+            active_blocks_count as f32 / total_capacity as f32
+        } else {
+            0.0
+        };
+
+        ForwardPassMetrics {
+            request_active_slots: state.running.len() as u64,
+            request_total_slots: 420, // Dummy value as specified
+            kv_active_blocks: active_blocks_count,
+            kv_total_blocks: total_capacity,
+            num_requests_waiting: state.waiting.len() as u64,
+            gpu_cache_usage_perc,
+            gpu_prefix_cache_hit_rate: 0.0, // Placeholder value as specified
+        }
+    }
+}
+
+/// Convert a Request to an ActiveSequence
+fn get_active_sequence(request: Request, block_size: usize, chunk_size: usize) -> ActiveSequence {
+    if let Request::Active(active_seq) = request {
+        return active_seq;
+    }
+
+    let Request::Direct(direct_request) = request else {
+        unreachable!("Request must be either Direct or Active");
+    };
+
+    ActiveSequence::new(
+        direct_request.tokens,
+        direct_request.max_output_tokens,
+        Some(block_size),
+        Some(chunk_size),
+    )
+}
+
+/// Processes MoveBlock signals with the KvManager.
+///
+/// When a signal fails, this function verifies that the failure is for an expected case:
+/// specifically a single signal attempting to create a single partial (generation) block.
+/// This validation is important because in normal operation, the only legitimate failure
+/// case should be when trying to acquire a new generation block - any other failures would
+/// indicate an unexpected state in the system.
+fn process_signals(
+    kv_manager_guard: &mut tokio::sync::MutexGuard<'_, KvManager>,
+    signals: &[MoveBlock],
+) -> bool {
+    for signal in signals {
+        if kv_manager_guard.process(signal) {
+            continue;
+        }
+
+        // Check we have a Use signal with blocks
+        let MoveBlock::Use(blocks, _) = signal else {
+            panic!("Failed signal is Invalid. Has to fail on generation signal.");
+        };
+
+        // Verify the signal contains exactly one block
+        if blocks.len() != 1 {
+            panic!("Failed signal is Invalid. Can have only one generation signal.");
+        }
+
+        // Verify the block is a PartialBlock (generation block)
+        if !matches!(blocks[0], UniqueBlock::PartialBlock(_)) {
+            panic!("Failed signal is Invalid. Generation block has to be partial.");
+        }
+
+        return false;
+    }
+
+    true
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use rstest::rstest;
+    use std::time::Duration;
+
+    #[rstest]
+    #[case::random(false)]
+    #[case::caching(true)]
+    #[tokio::test]
+    async fn test_scheduler_token_generation_patterns(#[case] use_shared_tokens: bool) {
+        std::env::set_var("RUST_LOG", "debug");
+
+        let kv_capacity: usize = 500;
+        let watermark: f64 = 0.01; // 1% watermark
+        let block_size: usize = 64;
+        let chunk_size: usize = 256;
+        let num_requests: usize = 100;
+        let input_len: usize = 1000;
+        let max_output_tokens: usize = 100;
+
+        // Create channel for token output
+        let (output_tx, mut output_rx) = mpsc::channel::<Uuid>(1024);
+
+        // Create scheduler with internal KvManager
+        let scheduler = Scheduler::new(
+            kv_capacity,
+            watermark,
+            block_size,
+            Some(chunk_size),
+            Some(output_tx),
+            None,
+        );
+
+        // Create shared tokens for caching case
+        let shared_tokens = if use_shared_tokens {
+            Some(
+                (0..input_len / 2)
+                    .map(|_| rand::random::<u32>() % 50000)
+                    .collect::<Vec<_>>(),
+            )
+        } else {
+            None
+        };
+
+        // Create test requests
+        for _ in 0..num_requests {
+            let input_tokens = if let Some(ref shared) = shared_tokens {
+                // For caching case: use shared tokens for first half, random for second half
+                let mut tokens = shared.clone();
+                tokens.extend((0..input_len / 2).map(|_| rand::random::<u32>() % 50000));
+                tokens
+            } else {
+                // For random case: create unique random token vector for each request
+                (0..input_len)
+                    .map(|_| rand::random::<u32>() % 50000)
+                    .collect::<Vec<_>>()
+            };
+
+            let request = DirectRequest {
+                tokens: input_tokens,
+                max_output_tokens,
+                uuid: None,
+            };
+            scheduler.receive(request).await;
+        }
+
+        let start_time = std::time::Instant::now();
+
+        // Collect all generated tokens (should be num_requests * max_output_tokens)
+        let expected_tokens = num_requests * max_output_tokens;
+        let mut received_tokens = 0;
+
+        // Set up a timeout that causes the test to panic if no tokens are received for 2 seconds
+        let timeout = tokio::time::sleep(Duration::from_secs(2));
+        tokio::pin!(timeout);
+
+        // Set up debug ticker interval
+        let mut debug_interval = interval(Duration::from_millis(500));
+
+        loop {
+            tokio::select! {
+                biased;
+
+                // Manual debug ticker that prints forward pass metrics
+                _ = debug_interval.tick() => {
+                    let _metrics = scheduler.get_forward_pass_metrics().await;
+                    // println!("Forward Pass Metrics: {:#?}", _metrics);
+                }
+
+                Some(_) = output_rx.recv() => {
+                    received_tokens += 1;
+                    // Reset timeout whenever we receive a token
+                    timeout.set(tokio::time::sleep(Duration::from_secs(2)));
+                }
+
+                _ = &mut timeout => {
+                    // Break instead of panicking when timeout occurs
+                    break;
+                }
+            }
+        }
+
+        // Calculate and print elapsed time
+        let elapsed = start_time.elapsed();
+        println!(
+            "Test completed in: {:?} for {} case",
+            elapsed,
+            if use_shared_tokens {
+                "caching"
+            } else {
+                "random"
+            }
+        );
+
+        // Assert that we received the expected number of tokens
+        assert!(
+            received_tokens > expected_tokens,
+            "Received {} tokens but expected more than {}",
+            received_tokens,
+            expected_tokens
+        );
+    }
+}
diff --git a/lib/llm/src/mocker/sequence.rs b/lib/llm/src/mocker/sequence.rs
new file mode 100644
index 000000000000..d53dd870e14e
--- /dev/null
+++ b/lib/llm/src/mocker/sequence.rs
@@ -0,0 +1,433 @@
+// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use crate::mocker::protocols::{MoveBlock, UniqueBlock};
+use crate::tokens::{TokenBlockSequence, Tokens};
+use derive_getters::Getters;
+use rand::random;
+use uuid;
+
+/// Create unique blocks from a TokenBlockSequence
+fn create_unique_blocks_from_sequence(
+    tokens: &TokenBlockSequence,
+    uuid: Option<uuid::Uuid>,
+    block_size: usize,
+) -> Vec<UniqueBlock> {
+    let mut unique_blocks: Vec<UniqueBlock> = tokens
+        .blocks()
+        .iter()
+        .map(|block| UniqueBlock::FullBlock(block.sequence_hash()))
+        .collect();
+
+    // Only push the partial block if tokens count isn't a multiple of block_size
+    if tokens.total_tokens() % block_size != 0 {
+        unique_blocks.push(match uuid {
+            Some(uuid) => UniqueBlock::PartialBlock(uuid),
+            None => UniqueBlock::default(),
+        });
+    }
+    unique_blocks
+}
+
+/// A sequence that is actively being built, with the ability to add tokens and commit to hashes
+/// TODO: reuse tokens
+#[derive(Debug, Getters)]
+pub struct ActiveSequence {
+    unique_blocks: Vec<UniqueBlock>,
+
+    tokens: TokenBlockSequence,
+
+    #[getter(copy)]
+    block_size: usize,
+
+    #[getter(copy)]
+    chunk_size: usize, // TODO: not actually used
+
+    #[getter(copy)]
+    max_output_tokens: usize,
+
+    #[getter(copy)]
+    generated_tokens: usize,
+
+    #[getter(copy)]
+    num_input_tokens: usize,
+
+    creation_signal: Option<MoveBlock>,
+}
+
+impl ActiveSequence {
+    /// Create a new ActiveSequence instance with the provided tokens
+    pub fn new(
+        tokens: Vec<u32>,
+        max_output_tokens: usize,
+        block_size: Option<usize>,
+        chunk_size: Option<usize>,
+    ) -> Self {
+        let block_size = block_size.unwrap_or(64);
+        assert!(block_size > 1, "block_size must be greater than 1");
+        let chunk_size = chunk_size.unwrap_or(256);
+        let num_input_tokens = tokens.len();
+
+        let tokens = Tokens::from(tokens).into_sequence(block_size, None);
+        let unique_blocks = create_unique_blocks_from_sequence(&tokens, None, block_size);
+        let creation_signal = Some(MoveBlock::Use(unique_blocks.clone(), None));
+
+        Self {
+            unique_blocks,
+            tokens,
+            block_size,
+            chunk_size,
+            max_output_tokens,
+            generated_tokens: 0,
+            num_input_tokens,
+            creation_signal,
+        }
+    }
+
+    pub fn extra_tokens(&self) -> usize {
+        self.len() % self.block_size
+    }
+
+    pub fn len(&self) -> usize {
+        self.tokens.total_tokens()
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.tokens.total_tokens() == 0
+    }
+
+    /// Create a new ActiveSequence instance and return the creation signal
+    pub fn new_with_signal(
+        tokens: Vec<u32>,
+        max_output_tokens: usize,
+        block_size: Option<usize>,
+        chunk_size: Option<usize>,
+    ) -> (Self, Option<MoveBlock>) {
+        let mut sequence = Self::new(tokens, max_output_tokens, block_size, chunk_size);
+        let signal = sequence.creation_signal.take();
+        (sequence, signal)
+    }
+
+    /// Push a token to the sequence
+    pub fn push(&mut self, token: u32) -> Option<Vec<MoveBlock>> {
+        self.tokens.append(token).expect("Token push failed.");
+        self.generated_tokens += 1;
+
+        if self.len() % self.block_size != 1 {
+            return None;
+        }
+
+        // Add a partial block for the first token in a new partial sequence
+        // Send Use signal (to allocate space for this new generation block)
+        let mut signals = Vec::new();
+
+        // Replace last partial block with full block if it exists
+        if let Some(UniqueBlock::PartialBlock(uuid)) = self.unique_blocks.last().cloned() {
+            let last_block_hash = self.tokens.last_complete_block().unwrap().sequence_hash();
+            self.unique_blocks.pop();
+            self.unique_blocks
+                .push(UniqueBlock::FullBlock(last_block_hash));
+            signals.push(MoveBlock::Promote(uuid, last_block_hash));
+        }
+
+        let new_partial_block = UniqueBlock::default();
+        self.unique_blocks.push(new_partial_block.clone());
+        signals.push(MoveBlock::Use(vec![new_partial_block], None));
+        Some(signals)
+    }
+
+    /// Generate a random token, push it to the sequence, and increment generation count.
+    ///
+    /// This function:
+    /// - Generates a random token and adds it to the current sequence
+    /// - Acquires a new partial block if needed or promotes an existing partial block to a full block
+    /// - Returns appropriate signals for the KvManager to process
+    ///
+    /// # Panics
+    ///
+    /// Calling this function when max_output_tokens has already been reached will cause a panic.
+    /// Always check `generated_tokens < max_output_tokens` before calling this method.
+    pub fn generate(&mut self) -> Vec<MoveBlock> {
+        // Assert that we haven't reached the maximum output tokens
+        assert!(
+            self.generated_tokens < self.max_output_tokens,
+            "Cannot generate more tokens: reached max_output_tokens limit"
+        );
+
+        // Generate a random token
+        let token = random::<u32>();
+
+        // Collect signals
+        let mut signals = Vec::new();
+
+        // Push the token to the sequence and collect any signals
+        if let Some(move_blocks) = self.push(token) {
+            signals.extend(move_blocks);
+        }
+
+        // Check if we've reached the limit after pushing
+        if self.generated_tokens != self.max_output_tokens {
+            return signals;
+        }
+
+        // Free all blocks when we reach max tokens
+        signals.extend(self.free_signal());
+        signals
+    }
+
+    /// Free all blocks, generating appropriate signals for each block type
+    pub fn free_signal(&self) -> Vec<MoveBlock> {
+        self.unique_blocks
+            .iter()
+            .rev()
+            .map(|block| match block {
+                UniqueBlock::PartialBlock(uuid) => {
+                    MoveBlock::Destroy(vec![UniqueBlock::PartialBlock(*uuid)])
+                }
+                UniqueBlock::FullBlock(hash) => {
+                    MoveBlock::Deref(vec![UniqueBlock::FullBlock(*hash)])
+                }
+            })
+            .collect()
+    }
+
+    /// Reset the sequence to its initial state and return the free signals from freeing current blocks
+    /// maintaining the uuid of the last partial block
+    pub fn reset_with_signal(&mut self) -> Vec<MoveBlock> {
+        let free_signal = self.free_signal();
+
+        self.tokens.truncate(self.num_input_tokens).unwrap();
+        self.unique_blocks =
+            create_unique_blocks_from_sequence(&self.tokens, None, self.block_size);
+        self.generated_tokens = 0;
+        self.creation_signal = Some(MoveBlock::Use(self.unique_blocks.clone(), None));
+
+        free_signal
+    }
+
+    /// Pops last token in the sequence.
+    pub fn pop(&mut self) {
+        self.tokens.pop();
+        self.generated_tokens = self.generated_tokens.saturating_sub(1);
+
+        // Reverts to the last full block
+        if self.tokens.total_tokens() % self.block_size == 0 {
+            self.unique_blocks.pop();
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_active_sequence_push() {
+        // Create a sequence with block size 16 initialized with tokens [0..15]
+        let initial_tokens: Vec<u32> = (0..15).collect();
+        let (mut seq1, signal1) =
+            ActiveSequence::new_with_signal(initial_tokens, 100, Some(16), Some(256));
+        assert_eq!(seq1.num_input_tokens(), 15);
+        assert_eq!(seq1.len(), 15);
+
+        // Check that we got a Use signal
+        assert!(signal1.is_some());
+        match &signal1 {
+            Some(MoveBlock::Use(blocks, _)) => {
+                assert_eq!(blocks.len(), 1);
+            }
+            _ => panic!("Expected Use signal"),
+        }
+
+        // Push token 15 which should complete the block (no signals yet)
+        let signal_15 = seq1.push(15);
+        assert!(
+            signal_15.is_none(),
+            "Completing a block should not trigger signals"
+        );
+
+        // Push token 16 which should trigger both Promote and Use signals
+        let signal_16 = seq1.push(16);
+        assert!(signal_16.is_some());
+        let signal_16 = signal_16.unwrap();
+        assert_eq!(signal_16.len(), 2);
+
+        // Second signal should be Use for new partial block
+        match &signal_16[1] {
+            MoveBlock::Use(blocks, _) => {
+                assert_eq!(blocks.len(), 1);
+                assert!(matches!(blocks[0], UniqueBlock::PartialBlock(_)));
+            }
+            _ => panic!("Expected Use signal as first signal"),
+        }
+
+        // First signal should be Promote for the previous block
+        match &signal_16[0] {
+            MoveBlock::Promote(uuid, _) => {
+                // The uuid is generated dynamically, so we just check it exists
+                let _ = uuid;
+            }
+            _ => panic!("Expected Promote signal as second signal"),
+        }
+
+        // Verify state after pushing tokens
+        assert_eq!(seq1.unique_blocks().len(), 2); // One full block and one partial block
+        assert_eq!(seq1.len(), 17);
+        assert_eq!(seq1.len() % seq1.block_size(), 1);
+
+        // Create another sequence with block size 16 initialized with tokens [0..17]
+        let extended_tokens: Vec<u32> = (0..16).collect();
+        let (mut seq2, _) =
+            ActiveSequence::new_with_signal(extended_tokens, 100, Some(16), Some(256));
+        seq2.push(16);
+        seq2.pop();
+        seq2.push(16);
+
+        // Simplified assertions
+        assert_eq!(
+            seq1.unique_blocks()[0],
+            seq2.unique_blocks()[0],
+            "First blocks should be the same"
+        );
+
+        assert_ne!(
+            seq1.unique_blocks()[1],
+            seq2.unique_blocks()[1],
+            "Second blocks should be different"
+        );
+
+        // Reset partial block on seq1 and push back token 16
+        seq1.push(17);
+        seq1.pop();
+        seq1.pop();
+        seq1.push(16);
+
+        // Now push tokens 17..32 to both sequences
+        for token in 17..33 {
+            seq1.push(token);
+            seq2.push(token);
+        }
+
+        // Both sequences should now have 2 blocks:
+        // 1. FullBlock for tokens 0-15
+        // 2. FullBlock for tokens 16-31
+        // 3. No partial block since there are no remaining tokens
+        assert_eq!(
+            seq1.unique_blocks().len(),
+            3,
+            "seq1 should have exactly 3 blocks"
+        );
+        assert_eq!(
+            seq2.unique_blocks().len(),
+            3,
+            "seq2 should have exactly 3 blocks"
+        );
+        assert_eq!(
+            seq1.len() % seq1.block_size(),
+            1,
+            "seq1 should have 1 partial token"
+        );
+        assert_eq!(
+            seq2.len() % seq2.block_size(),
+            1,
+            "seq2 should have 1 partial token"
+        );
+
+        // Verify that both sequences have identical blocks up to the second position
+        assert_eq!(
+            &seq1.unique_blocks()[0..2],
+            &seq2.unique_blocks()[0..2],
+            "First two blocks should be identical"
+        );
+
+        // Reset seq1 and check that it equals the original clone
+        let free_signals = seq1.reset_with_signal();
+
+        // Verify the reset signals include proper cleanup events
+        assert!(!free_signals.is_empty());
+    }
+
+    #[test]
+    fn test_active_sequence_generate_signals() {
+        // Create a sequence with block size 16, max_output_tokens 4, initialized with tokens [0..14)
+        let initial_tokens: Vec<u32> = (0..14).collect();
+        let (mut seq, signal) =
+            ActiveSequence::new_with_signal(initial_tokens, 5, Some(16), Some(256));
+
+        // Initial signal - should have received a Use signal for the partial block
+        assert!(signal.is_some());
+        match signal {
+            Some(MoveBlock::Use(blocks, _)) => {
+                assert_eq!(blocks.len(), 1);
+                assert!(matches!(blocks[0], UniqueBlock::PartialBlock(_)));
+            }
+            _ => panic!("Expected Use signal for the initial partial block"),
+        }
+
+        // Generate first two tokens - should not trigger new signals
+        seq.generate();
+        let signals_first = seq.generate();
+        assert_eq!(signals_first.len(), 0);
+
+        // Generate third token - this fills the block and should trigger both Promote and Use signals
+        let signals_second = seq.generate();
+        assert_eq!(signals_second.len(), 2);
+
+        // First signal should be Use for new partial block
+        match &signals_second[1] {
+            MoveBlock::Use(blocks, _) => {
+                assert_eq!(blocks.len(), 1);
+                assert!(matches!(blocks[0], UniqueBlock::PartialBlock(_)));
+            }
+            _ => panic!("Expected Use signal as second signal after second token"),
+        }
+
+        // Second signal should be Promote
+        match &signals_second[0] {
+            MoveBlock::Promote(uuid, hash) => {
+                // The uuid and hash values are generated dynamically, so we just check the event type
+                let _ = uuid;
+                let _ = hash;
+            }
+            _ => panic!("Expected Promote signal as first signal after second token"),
+        }
+
+        // Generate fourth token - should not trigger new signals as it's adding to partial block
+        let signals_third = seq.generate();
+        assert_eq!(signals_third.len(), 0);
+
+        // Generate last token - we reach max_output_tokens, should trigger Destroy and Deref signals
+        let signals_last = seq.generate();
+        assert_eq!(signals_last.len(), 2);
+
+        // First signal should be Destroy for the partial block
+        match &signals_last[0] {
+            MoveBlock::Destroy(blocks) => {
+                assert_eq!(blocks.len(), 1);
+                assert!(matches!(blocks[0], UniqueBlock::PartialBlock(_)));
+            }
+            _ => panic!("Expected Destroy signal for partial block after fourth token"),
+        }
+
+        // Second signal should be Deref for the full block
+        match &signals_last[1] {
+            MoveBlock::Deref(blocks) => {
+                assert_eq!(blocks.len(), 1);
+                assert!(matches!(blocks[0], UniqueBlock::FullBlock(_)));
+            }
+            _ => panic!("Expected Deref signal for full block after fourth token"),
+        }
+    }
+}
diff --git a/lib/llm/src/tokens.rs b/lib/llm/src/tokens.rs
index 92728a25a246..645919d1bdb5 100644
--- a/lib/llm/src/tokens.rs
+++ b/lib/llm/src/tokens.rs
@@ -188,7 +188,7 @@ pub enum TokenBlockError {
 ///
 /// This structure accumulates tokens until it reaches the specified `block_size`,
 /// at which point it can be [`commit`](PartialTokenBlock::commit)ted into a full [`TokenBlock`].
-#[derive(Debug)] // No Clone: intended to be unique within a sequence
+#[derive(Debug, PartialEq)] // No Clone: intended to be unique within a sequence
 pub struct PartialTokenBlock {
     tokens: Tokens,
     block_size: usize,
@@ -478,7 +478,7 @@ impl TokenBlock {
 /// - [`BlockHash`]: Hash of tokens within a single block (seeded by [`SaltHash`]).
 /// - [`SequenceHash`]: Hash combining the previous block's [`SequenceHash`] and the current
 ///   block's [`BlockHash`] (also seeded by [`SaltHash`]).
-#[derive(Debug)]
+#[derive(Debug, PartialEq)]
 pub struct TokenBlockSequence {
     blocks: Vec<TokenBlock>,
     current_block: PartialTokenBlock,


Reviews:
user: login: jthomson04
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-19T20:14:42Z
user: login: PeaBrane
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-19T20:40:09Z
user: login: PeaBrane
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-19T20:44:04Z
user: login: PeaBrane
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-19T20:50:18Z
user: login: PeaBrane
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-19T20:53:19Z
user: login: PeaBrane
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-19T21:08:15Z
user: login: PeaBrane
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-19T21:47:46Z
user: login: PeaBrane
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-19T22:11:46Z
user: login: PeaBrane
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-19T22:49:57Z
user: login: PeaBrane
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-19T22:52:05Z
user: login: alec-flowers
body: Can you add headers to each file describing the purpose and goal? You did a great job in the PR description. It would be good to translate that into the code. 

Also I think would be useful once its set up to look at what sort of numbers can be generated by running these workers. 

We may want to have it where the Block Manager actually emits Events that the KVRouter can receive the signal and do something with. This would mean adding functionality from the KVPublisher to the MockWorker.

state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-21T00:30:19Z

Review comments:
diff_hunk: @@ -0,0 +1,396 @@
+// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
path: lib/llm/src/mocker/kv_manager.rs
user: login: jthomson04
body: Moving forward, it would be nice to run the actual block manager here. Could allocate just a host pool with a super small block size to minimize overhead.
created_at: 2025-05-19T18:16:56Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,396 @@
+// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use crate::mocker::evictor::LRUEvictor;
+use crate::mocker::protocols::{MoveBlock, MoveBlockResponse, PrefillCost, UniqueBlock};
+use crate::mocker::sequence::ActiveSequence;
+use std::collections::{HashMap, HashSet};
+
+/// Mock implementation of worker for testing and simulation
+pub struct KvManager {
+    max_capacity: usize,
+    block_size: usize,
+    active_blocks: HashMap<UniqueBlock, usize>,
+    inactive_blocks: LRUEvictor<UniqueBlock>,
+    all_blocks: HashSet<UniqueBlock>,
+}
+
+impl KvManager {
+    pub fn new(max_capacity: usize, block_size: usize) -> Self {
+        let active_blocks = HashMap::new();
+        let inactive_blocks = LRUEvictor::new();
+        let all_blocks = HashSet::new();
+
+        KvManager {
+            max_capacity,
+            block_size,
+            active_blocks,
+            inactive_blocks,
+            all_blocks,
+        }
+    }
+
+    /// Get the maximum capacity
+    pub fn max_capacity(&self) -> usize {
path: lib/llm/src/mocker/kv_manager.rs
user: login: jthomson04
body: Could use derive_getter here (https://docs.rs/derive-getters/latest/derive_getters/)
See example usage here: https://github.com/ai-dynamo/dynamo/blob/eb133e3fa054ae52cc901a8af40a720ec6de1b23/lib/tokens/src/lib.rs#L220
created_at: 2025-05-19T18:18:31Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,193 @@
+// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::cmp::Eq;
+use std::collections::{HashMap, VecDeque};
+use std::hash::Hash;
+use std::time::Instant;
+
+/// An LRU evictor that maintains objects and evicts them based on their
+/// last accessed time. Implements a "lazy" eviction mechanism where:
+/// 1. The priority queue does not immediately reflect updates or removes
+/// 2. Objects are pushed to the queue in order of increasing priority (older objects first)
+/// 3. The user must ensure objects are added in correct priority (temporal order)
+/// 4. Remove and update operations are lazy - entries remain in the queue until
+///    they are either evicted or cleaned up during maintenance
+#[derive(Debug)]
+pub struct LRUEvictor<T: Clone + Eq + Hash> {
+    free_table: HashMap<T, f64>,
+    priority_queue: VecDeque<(T, f64)>,
+    cleanup_threshold: usize,
+    start_time: Instant,
+}
+
+impl<T: Clone + Eq + Hash> Default for LRUEvictor<T> {
path: lib/llm/src/mocker/evictor.rs
user: login: jthomson04
body: Unless we're planning on adding a bunch of arguments, this default trait probably isn't required
created_at: 2025-05-19T18:32:08Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,396 @@
+// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use crate::mocker::evictor::LRUEvictor;
+use crate::mocker::protocols::{MoveBlock, MoveBlockResponse, PrefillCost, UniqueBlock};
+use crate::mocker::sequence::ActiveSequence;
+use std::collections::{HashMap, HashSet};
+
+/// Mock implementation of worker for testing and simulation
+pub struct KvManager {
+    max_capacity: usize,
+    block_size: usize,
+    active_blocks: HashMap<UniqueBlock, usize>,
+    inactive_blocks: LRUEvictor<UniqueBlock>,
+    all_blocks: HashSet<UniqueBlock>,
+}
+
+impl KvManager {
+    pub fn new(max_capacity: usize, block_size: usize) -> Self {
+        let active_blocks = HashMap::new();
+        let inactive_blocks = LRUEvictor::new();
+        let all_blocks = HashSet::new();
+
+        KvManager {
+            max_capacity,
+            block_size,
+            active_blocks,
+            inactive_blocks,
+            all_blocks,
+        }
+    }
+
+    /// Get the maximum capacity
+    pub fn max_capacity(&self) -> usize {
+        self.max_capacity
+    }
+
+    /// Get the block size
+    pub fn block_size(&self) -> usize {
+        self.block_size
+    }
+
+    /// Get a reference to the active blocks
+    pub fn active_blocks(&self) -> &HashMap<UniqueBlock, usize> {
+        &self.active_blocks
+    }
+
+    /// Get a reference to the inactive blocks
+    pub fn inactive_blocks(&self) -> &LRUEvictor<UniqueBlock> {
+        &self.inactive_blocks
+    }
+
+    /// Get a reference to all blocks
+    pub fn all_blocks(&self) -> &HashSet<UniqueBlock> {
+        &self.all_blocks
+    }
+
+    /// Process a MoveBlock instruction synchronously
+    pub fn process(&mut self, event: &MoveBlock) -> MoveBlockResponse {
+        match event {
+            MoveBlock::Use(hashes, _) => {
+                for hash in hashes {
+                    // First check if it already exists in active blocks
+                    if let Some(ref_count) = self.active_blocks.get_mut(hash) {
+                        // Block already active, just increment reference count
+                        *ref_count += 1;
+                        continue;
+                    }
+
+                    // Then check if it exists in inactive and move it to active if found
+                    if self.inactive_blocks.remove(hash) {
+                        // Insert into active with reference count 1
+                        self.active_blocks.insert(hash.clone(), 1);
+                        continue;
+                    }
+
+                    // Get counts for capacity check
+                    let active_count = self.active_blocks.len();
+                    let inactive_count = self.inactive_blocks.len();
+
+                    // If at max capacity, evict the oldest entry from inactive blocks
+                    if active_count + inactive_count >= self.max_capacity {
+                        if let Some(evicted) = self.inactive_blocks.evict() {
+                            // Remove evicted block from all_blocks
+                            self.all_blocks.remove(&evicted);
+                        } else {
+                            // Return failure instead of panicking
+                            return MoveBlockResponse::Failure;
+                        }
+                    }
+
+                    // Now insert the new block in active blocks with reference count 1
+                    self.active_blocks.insert(hash.clone(), 1);
+                    // Add to all_blocks as it's a new block
+                    self.all_blocks.insert(hash.clone());
+                }
+            }
+            MoveBlock::Destroy(hashes) => {
+                // Loop in inverse direction
+                for hash in hashes.iter().rev() {
+                    self.active_blocks.remove(hash).unwrap();
+                    // Remove from all_blocks when destroyed
+                    assert!(self.all_blocks.remove(hash));
+                }
+            }
+            MoveBlock::Deref(hashes) => {
+                // Loop in inverse direction
+                for hash in hashes.iter().rev() {
path: lib/llm/src/mocker/kv_manager.rs
user: login: jthomson04
body: Instead of manually tracking reference counts, an alternative could be to use `Arc` with a custom `Drop` implementation. That's sortof similar to what's currently done in the kvbm.
created_at: 2025-05-19T18:40:08Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,396 @@
+// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use crate::mocker::evictor::LRUEvictor;
+use crate::mocker::protocols::{MoveBlock, MoveBlockResponse, PrefillCost, UniqueBlock};
+use crate::mocker::sequence::ActiveSequence;
+use std::collections::{HashMap, HashSet};
+
+/// Mock implementation of worker for testing and simulation
+pub struct KvManager {
+    max_capacity: usize,
+    block_size: usize,
+    active_blocks: HashMap<UniqueBlock, usize>,
+    inactive_blocks: LRUEvictor<UniqueBlock>,
+    all_blocks: HashSet<UniqueBlock>,
+}
+
+impl KvManager {
+    pub fn new(max_capacity: usize, block_size: usize) -> Self {
+        let active_blocks = HashMap::new();
+        let inactive_blocks = LRUEvictor::new();
+        let all_blocks = HashSet::new();
+
+        KvManager {
+            max_capacity,
+            block_size,
+            active_blocks,
+            inactive_blocks,
+            all_blocks,
+        }
+    }
+
+    /// Get the maximum capacity
+    pub fn max_capacity(&self) -> usize {
+        self.max_capacity
+    }
+
+    /// Get the block size
+    pub fn block_size(&self) -> usize {
+        self.block_size
+    }
+
+    /// Get a reference to the active blocks
+    pub fn active_blocks(&self) -> &HashMap<UniqueBlock, usize> {
+        &self.active_blocks
+    }
+
+    /// Get a reference to the inactive blocks
+    pub fn inactive_blocks(&self) -> &LRUEvictor<UniqueBlock> {
+        &self.inactive_blocks
+    }
+
+    /// Get a reference to all blocks
+    pub fn all_blocks(&self) -> &HashSet<UniqueBlock> {
+        &self.all_blocks
+    }
+
+    /// Process a MoveBlock instruction synchronously
+    pub fn process(&mut self, event: &MoveBlock) -> MoveBlockResponse {
+        match event {
+            MoveBlock::Use(hashes, _) => {
+                for hash in hashes {
+                    // First check if it already exists in active blocks
+                    if let Some(ref_count) = self.active_blocks.get_mut(hash) {
+                        // Block already active, just increment reference count
+                        *ref_count += 1;
+                        continue;
+                    }
+
+                    // Then check if it exists in inactive and move it to active if found
+                    if self.inactive_blocks.remove(hash) {
+                        // Insert into active with reference count 1
+                        self.active_blocks.insert(hash.clone(), 1);
+                        continue;
+                    }
+
+                    // Get counts for capacity check
+                    let active_count = self.active_blocks.len();
+                    let inactive_count = self.inactive_blocks.len();
+
+                    // If at max capacity, evict the oldest entry from inactive blocks
+                    if active_count + inactive_count >= self.max_capacity {
+                        if let Some(evicted) = self.inactive_blocks.evict() {
+                            // Remove evicted block from all_blocks
+                            self.all_blocks.remove(&evicted);
+                        } else {
+                            // Return failure instead of panicking
+                            return MoveBlockResponse::Failure;
+                        }
+                    }
+
+                    // Now insert the new block in active blocks with reference count 1
+                    self.active_blocks.insert(hash.clone(), 1);
+                    // Add to all_blocks as it's a new block
+                    self.all_blocks.insert(hash.clone());
+                }
+            }
+            MoveBlock::Destroy(hashes) => {
+                // Loop in inverse direction
+                for hash in hashes.iter().rev() {
+                    self.active_blocks.remove(hash).unwrap();
+                    // Remove from all_blocks when destroyed
+                    assert!(self.all_blocks.remove(hash));
+                }
+            }
+            MoveBlock::Deref(hashes) => {
+                // Loop in inverse direction
+                for hash in hashes.iter().rev() {
+                    // Decrement reference count and check if we need to move to inactive
+                    if let Some(ref_count) = self.active_blocks.get_mut(hash) {
+                        *ref_count -= 1;
+
+                        // If reference count reaches zero, remove from active and move to inactive
+                        if *ref_count == 0 {
+                            self.active_blocks.remove(hash);
+                            // Use the LRUEvictor's timing functionality
+                            self.inactive_blocks.insert(hash.clone());
+                        }
+                    }
+                }
+            }
+            MoveBlock::Promote(uuid, hash) => {
+                let uuid_block = UniqueBlock::PartialBlock(*uuid);
path: lib/llm/src/mocker/kv_manager.rs
user: login: jthomson04
body: Is there any way we can use the stuff in `crate::tokens` instead?
created_at: 2025-05-19T18:42:28Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,100 @@
+// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use serde::{Deserialize, Serialize};
+use uuid::Uuid;
+
+pub type Token = u32;
+pub type LocalBlockHash = u64;
+/// A global hash identifier for blocks
+pub type GlobalHash = u64;
+pub type NumBlocks = usize;
+
+/// Represents an active block in the cache with a reference count
+#[derive(Debug, Clone, Hash, Eq, PartialEq, Serialize, Deserialize)]
+pub enum UniqueBlock {
+    /// Block identified by UUID
+    PartialBlock(Uuid),
+    /// Block identified by hash
+    FullBlock(GlobalHash),
+}
+
+impl Default for UniqueBlock {
+    fn default() -> Self {
+        // Generate a random UUID when default is used
+        Self::PartialBlock(Uuid::new_v4())
+    }
+}
+
+/// Represents different block movement operations in the cache
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+pub enum MoveBlock {
+    Use(Vec<UniqueBlock>, Option<f64>),
+    Destroy(Vec<UniqueBlock>),
+    Deref(Vec<UniqueBlock>),
+    Promote(Uuid, GlobalHash),
+}
+
+/// Response to block movement operations in the cache
+#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
+pub enum MoveBlockResponse {
path: lib/llm/src/mocker/protocols.rs
user: login: jthomson04
body: Instead of this custom struct, just use a `Result<()>`
created_at: 2025-05-19T18:43:04Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,396 @@
+// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use crate::mocker::evictor::LRUEvictor;
+use crate::mocker::protocols::{MoveBlock, MoveBlockResponse, PrefillCost, UniqueBlock};
+use crate::mocker::sequence::ActiveSequence;
+use std::collections::{HashMap, HashSet};
+
+/// Mock implementation of worker for testing and simulation
+pub struct KvManager {
+    max_capacity: usize,
+    block_size: usize,
+    active_blocks: HashMap<UniqueBlock, usize>,
+    inactive_blocks: LRUEvictor<UniqueBlock>,
+    all_blocks: HashSet<UniqueBlock>,
+}
+
+impl KvManager {
+    pub fn new(max_capacity: usize, block_size: usize) -> Self {
+        let active_blocks = HashMap::new();
+        let inactive_blocks = LRUEvictor::new();
+        let all_blocks = HashSet::new();
+
+        KvManager {
+            max_capacity,
+            block_size,
+            active_blocks,
+            inactive_blocks,
+            all_blocks,
+        }
+    }
+
+    /// Get the maximum capacity
+    pub fn max_capacity(&self) -> usize {
+        self.max_capacity
+    }
+
+    /// Get the block size
+    pub fn block_size(&self) -> usize {
+        self.block_size
+    }
+
+    /// Get a reference to the active blocks
+    pub fn active_blocks(&self) -> &HashMap<UniqueBlock, usize> {
+        &self.active_blocks
+    }
+
+    /// Get a reference to the inactive blocks
+    pub fn inactive_blocks(&self) -> &LRUEvictor<UniqueBlock> {
+        &self.inactive_blocks
+    }
+
+    /// Get a reference to all blocks
+    pub fn all_blocks(&self) -> &HashSet<UniqueBlock> {
+        &self.all_blocks
+    }
+
+    /// Process a MoveBlock instruction synchronously
+    pub fn process(&mut self, event: &MoveBlock) -> MoveBlockResponse {
+        match event {
+            MoveBlock::Use(hashes, _) => {
+                for hash in hashes {
+                    // First check if it already exists in active blocks
+                    if let Some(ref_count) = self.active_blocks.get_mut(hash) {
+                        // Block already active, just increment reference count
+                        *ref_count += 1;
+                        continue;
+                    }
+
+                    // Then check if it exists in inactive and move it to active if found
+                    if self.inactive_blocks.remove(hash) {
+                        // Insert into active with reference count 1
+                        self.active_blocks.insert(hash.clone(), 1);
+                        continue;
+                    }
+
+                    // Get counts for capacity check
+                    let active_count = self.active_blocks.len();
+                    let inactive_count = self.inactive_blocks.len();
+
+                    // If at max capacity, evict the oldest entry from inactive blocks
+                    if active_count + inactive_count >= self.max_capacity {
+                        if let Some(evicted) = self.inactive_blocks.evict() {
+                            // Remove evicted block from all_blocks
+                            self.all_blocks.remove(&evicted);
+                        } else {
+                            // Return failure instead of panicking
+                            return MoveBlockResponse::Failure;
+                        }
+                    }
+
+                    // Now insert the new block in active blocks with reference count 1
+                    self.active_blocks.insert(hash.clone(), 1);
+                    // Add to all_blocks as it's a new block
+                    self.all_blocks.insert(hash.clone());
+                }
+            }
+            MoveBlock::Destroy(hashes) => {
+                // Loop in inverse direction
+                for hash in hashes.iter().rev() {
+                    self.active_blocks.remove(hash).unwrap();
+                    // Remove from all_blocks when destroyed
+                    assert!(self.all_blocks.remove(hash));
+                }
+            }
+            MoveBlock::Deref(hashes) => {
+                // Loop in inverse direction
+                for hash in hashes.iter().rev() {
+                    // Decrement reference count and check if we need to move to inactive
+                    if let Some(ref_count) = self.active_blocks.get_mut(hash) {
+                        *ref_count -= 1;
+
+                        // If reference count reaches zero, remove from active and move to inactive
+                        if *ref_count == 0 {
+                            self.active_blocks.remove(hash);
+                            // Use the LRUEvictor's timing functionality
+                            self.inactive_blocks.insert(hash.clone());
+                        }
+                    }
+                }
+            }
+            MoveBlock::Promote(uuid, hash) => {
+                let uuid_block = UniqueBlock::PartialBlock(*uuid);
+                let hash_block = UniqueBlock::FullBlock(*hash);
+
+                let Some(ref_count) = self.active_blocks.remove(&uuid_block) else {
+                    panic!("Cannot promote non-existent block.");
+                };
+
+                // Replace with hash block, keeping the same reference count
+                self.active_blocks.insert(hash_block.clone(), ref_count);
+
+                // Update all_blocks
+                assert!(self.all_blocks.remove(&uuid_block));
+                self.all_blocks.insert(hash_block);
+            }
+        }
+
+        // Return success if we made it this far
+        MoveBlockResponse::Success
+    }
+
+    /// Get the count of blocks in the input list that aren't in all_blocks
+    pub fn probe_new_blocks(&self, blocks: &[UniqueBlock]) -> usize {
+        blocks
+            .iter()
+            .filter(|&block| !self.all_blocks.contains(block))
+            .count()
+    }
+
+    /// Get the current capacity (active blocks + inactive blocks)
+    pub fn current_capacity(&self) -> usize {
+        let active = self.active_blocks.len();
+        let inactive = self.inactive_blocks.len();
+        active + inactive
+    }
+
+    /// Get the current capacity as a percentage of the maximum capacity
+    pub fn current_capacity_perc(&self) -> f64 {
+        let current = self.current_capacity() as f64;
+        current / self.max_capacity as f64
+    }
+
+    /// Get the keys of inactive blocks
+    pub fn get_inactive_blocks(&self) -> Vec<&UniqueBlock> {
+        self.inactive_blocks.keys().collect()
+    }
+
+    /// Get the keys of active blocks
+    pub fn get_active_blocks(&self) -> Vec<&UniqueBlock> {
+        self.active_blocks.keys().collect()
+    }
+
+    /// Check if a sequence can be scheduled and calculate cost if possible
+    pub fn try_schedule(
+        &self,
+        sequence: &ActiveSequence,
+        watermark: f64,
+        tokens_budget: usize,
+    ) -> Option<PrefillCost> {
+        // Return None immediately if tokens_budget is 0
+        if tokens_budget == 0 {
+            return None;
+        }
+
+        // Get unique blocks from the sequence
+        let unique_blocks = sequence.unique_blocks();
+
+        // Get the count of new blocks
+        let new_blocks = self.probe_new_blocks(unique_blocks);
+
+        // Calculate current usage and available capacity
+        let active_count = self.active_blocks.len();
+
+        // Check if we can schedule based on the watermark
+        if (active_count + new_blocks) as f64 > (1.0 - watermark) * self.max_capacity as f64 {
+            return None;
+        }
+
+        // Calculate overlap blocks
+        let overlap_blocks = unique_blocks.len() - new_blocks;
+
+        // Calculate new tokens
+        let new_tokens = sequence.num_input_tokens() - overlap_blocks * self.block_size;
+
+        // // Print the full equation with actual values substituted
+        // println!("{} = {} - ({} * {}) (new_tokens = num_input_tokens - overlap_blocks * block_size)",
+        //     new_tokens,
+        //     sequence.num_input_tokens(),
+        //     overlap_blocks,
+        //     self.block_size);
+
+        // Return None if new_tokens exceeds tokens_budget
+        if new_tokens > tokens_budget {
+            return None;
+        }
+
+        // Calculate prefill compute
+        let prefill_compute =
+            new_tokens as f64 * (new_tokens + overlap_blocks * self.block_size) as f64;
+
+        Some(PrefillCost {
+            new_tokens,
+            prefill_compute,
+        })
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_failure_on_max_capacity() {
+        // Create a KvManager with 10 blocks capacity
+        let mut manager = KvManager::new(10, 16);
+
+        // Helper function to use multiple blocks that returns the response
+        fn use_blocks(manager: &mut KvManager, ids: Vec<u64>) -> MoveBlockResponse {
+            let blocks = ids.into_iter().map(UniqueBlock::FullBlock).collect();
+            manager.process(&MoveBlock::Use(blocks, None))
+        }
+
+        // First use 10 blocks (0 to 9) in a batch
+        let response = use_blocks(&mut manager, (0..10).collect());
+        assert_eq!(response, MoveBlockResponse::Success);
path: lib/llm/src/mocker/kv_manager.rs
user: login: jthomson04
body: Use `Result<()>` as the return value here, and check `response.is_ok()`
created_at: 2025-05-19T19:51:00Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,396 @@
+// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use crate::mocker::evictor::LRUEvictor;
+use crate::mocker::protocols::{MoveBlock, MoveBlockResponse, PrefillCost, UniqueBlock};
+use crate::mocker::sequence::ActiveSequence;
+use std::collections::{HashMap, HashSet};
+
+/// Mock implementation of worker for testing and simulation
+pub struct KvManager {
+    max_capacity: usize,
+    block_size: usize,
+    active_blocks: HashMap<UniqueBlock, usize>,
+    inactive_blocks: LRUEvictor<UniqueBlock>,
+    all_blocks: HashSet<UniqueBlock>,
+}
+
+impl KvManager {
+    pub fn new(max_capacity: usize, block_size: usize) -> Self {
+        let active_blocks = HashMap::new();
+        let inactive_blocks = LRUEvictor::new();
+        let all_blocks = HashSet::new();
+
+        KvManager {
+            max_capacity,
+            block_size,
+            active_blocks,
+            inactive_blocks,
+            all_blocks,
+        }
+    }
+
+    /// Get the maximum capacity
+    pub fn max_capacity(&self) -> usize {
+        self.max_capacity
+    }
+
+    /// Get the block size
+    pub fn block_size(&self) -> usize {
+        self.block_size
+    }
+
+    /// Get a reference to the active blocks
+    pub fn active_blocks(&self) -> &HashMap<UniqueBlock, usize> {
+        &self.active_blocks
+    }
+
+    /// Get a reference to the inactive blocks
+    pub fn inactive_blocks(&self) -> &LRUEvictor<UniqueBlock> {
+        &self.inactive_blocks
+    }
+
+    /// Get a reference to all blocks
+    pub fn all_blocks(&self) -> &HashSet<UniqueBlock> {
+        &self.all_blocks
+    }
+
+    /// Process a MoveBlock instruction synchronously
+    pub fn process(&mut self, event: &MoveBlock) -> MoveBlockResponse {
+        match event {
+            MoveBlock::Use(hashes, _) => {
+                for hash in hashes {
+                    // First check if it already exists in active blocks
+                    if let Some(ref_count) = self.active_blocks.get_mut(hash) {
+                        // Block already active, just increment reference count
+                        *ref_count += 1;
+                        continue;
+                    }
+
+                    // Then check if it exists in inactive and move it to active if found
+                    if self.inactive_blocks.remove(hash) {
+                        // Insert into active with reference count 1
+                        self.active_blocks.insert(hash.clone(), 1);
+                        continue;
+                    }
+
+                    // Get counts for capacity check
+                    let active_count = self.active_blocks.len();
+                    let inactive_count = self.inactive_blocks.len();
+
+                    // If at max capacity, evict the oldest entry from inactive blocks
+                    if active_count + inactive_count >= self.max_capacity {
+                        if let Some(evicted) = self.inactive_blocks.evict() {
+                            // Remove evicted block from all_blocks
+                            self.all_blocks.remove(&evicted);
+                        } else {
+                            // Return failure instead of panicking
+                            return MoveBlockResponse::Failure;
+                        }
+                    }
+
+                    // Now insert the new block in active blocks with reference count 1
+                    self.active_blocks.insert(hash.clone(), 1);
+                    // Add to all_blocks as it's a new block
+                    self.all_blocks.insert(hash.clone());
+                }
+            }
+            MoveBlock::Destroy(hashes) => {
+                // Loop in inverse direction
+                for hash in hashes.iter().rev() {
+                    self.active_blocks.remove(hash).unwrap();
+                    // Remove from all_blocks when destroyed
+                    assert!(self.all_blocks.remove(hash));
+                }
+            }
+            MoveBlock::Deref(hashes) => {
+                // Loop in inverse direction
+                for hash in hashes.iter().rev() {
+                    // Decrement reference count and check if we need to move to inactive
+                    if let Some(ref_count) = self.active_blocks.get_mut(hash) {
+                        *ref_count -= 1;
+
+                        // If reference count reaches zero, remove from active and move to inactive
+                        if *ref_count == 0 {
+                            self.active_blocks.remove(hash);
+                            // Use the LRUEvictor's timing functionality
+                            self.inactive_blocks.insert(hash.clone());
+                        }
+                    }
+                }
+            }
+            MoveBlock::Promote(uuid, hash) => {
+                let uuid_block = UniqueBlock::PartialBlock(*uuid);
+                let hash_block = UniqueBlock::FullBlock(*hash);
+
+                let Some(ref_count) = self.active_blocks.remove(&uuid_block) else {
+                    panic!("Cannot promote non-existent block.");
+                };
+
+                // Replace with hash block, keeping the same reference count
+                self.active_blocks.insert(hash_block.clone(), ref_count);
+
+                // Update all_blocks
+                assert!(self.all_blocks.remove(&uuid_block));
+                self.all_blocks.insert(hash_block);
+            }
+        }
+
+        // Return success if we made it this far
+        MoveBlockResponse::Success
+    }
+
+    /// Get the count of blocks in the input list that aren't in all_blocks
+    pub fn probe_new_blocks(&self, blocks: &[UniqueBlock]) -> usize {
+        blocks
+            .iter()
+            .filter(|&block| !self.all_blocks.contains(block))
+            .count()
+    }
+
+    /// Get the current capacity (active blocks + inactive blocks)
+    pub fn current_capacity(&self) -> usize {
+        let active = self.active_blocks.len();
+        let inactive = self.inactive_blocks.len();
+        active + inactive
+    }
+
+    /// Get the current capacity as a percentage of the maximum capacity
+    pub fn current_capacity_perc(&self) -> f64 {
+        let current = self.current_capacity() as f64;
+        current / self.max_capacity as f64
+    }
+
+    /// Get the keys of inactive blocks
+    pub fn get_inactive_blocks(&self) -> Vec<&UniqueBlock> {
+        self.inactive_blocks.keys().collect()
+    }
+
+    /// Get the keys of active blocks
+    pub fn get_active_blocks(&self) -> Vec<&UniqueBlock> {
+        self.active_blocks.keys().collect()
+    }
+
+    /// Check if a sequence can be scheduled and calculate cost if possible
+    pub fn try_schedule(
+        &self,
+        sequence: &ActiveSequence,
+        watermark: f64,
+        tokens_budget: usize,
+    ) -> Option<PrefillCost> {
+        // Return None immediately if tokens_budget is 0
+        if tokens_budget == 0 {
+            return None;
+        }
+
+        // Get unique blocks from the sequence
+        let unique_blocks = sequence.unique_blocks();
+
+        // Get the count of new blocks
+        let new_blocks = self.probe_new_blocks(unique_blocks);
+
+        // Calculate current usage and available capacity
+        let active_count = self.active_blocks.len();
+
+        // Check if we can schedule based on the watermark
+        if (active_count + new_blocks) as f64 > (1.0 - watermark) * self.max_capacity as f64 {
+            return None;
+        }
+
+        // Calculate overlap blocks
+        let overlap_blocks = unique_blocks.len() - new_blocks;
+
+        // Calculate new tokens
+        let new_tokens = sequence.num_input_tokens() - overlap_blocks * self.block_size;
+
+        // // Print the full equation with actual values substituted
+        // println!("{} = {} - ({} * {}) (new_tokens = num_input_tokens - overlap_blocks * block_size)",
+        //     new_tokens,
+        //     sequence.num_input_tokens(),
+        //     overlap_blocks,
+        //     self.block_size);
+
+        // Return None if new_tokens exceeds tokens_budget
+        if new_tokens > tokens_budget {
+            return None;
+        }
+
+        // Calculate prefill compute
+        let prefill_compute =
+            new_tokens as f64 * (new_tokens + overlap_blocks * self.block_size) as f64;
+
+        Some(PrefillCost {
+            new_tokens,
+            prefill_compute,
+        })
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_failure_on_max_capacity() {
+        // Create a KvManager with 10 blocks capacity
+        let mut manager = KvManager::new(10, 16);
+
+        // Helper function to use multiple blocks that returns the response
+        fn use_blocks(manager: &mut KvManager, ids: Vec<u64>) -> MoveBlockResponse {
+            let blocks = ids.into_iter().map(UniqueBlock::FullBlock).collect();
+            manager.process(&MoveBlock::Use(blocks, None))
+        }
+
+        // First use 10 blocks (0 to 9) in a batch
+        let response = use_blocks(&mut manager, (0..10).collect());
+        assert_eq!(response, MoveBlockResponse::Success);
+
+        // Verify we are at capacity
+        assert_eq!(manager.current_capacity(), 10);
+
+        // The 11th block should return Failure, not panic
+        let response = use_blocks(&mut manager, vec![10]);
+        assert_eq!(
+            response,
+            MoveBlockResponse::Failure,
+            "Expected Failure response when exceeding max capacity"
+        );
+    }
+
+    #[test]
+    // This is taken directly from the example in the vllm v1 prefix caching docs
+    fn test_block_lifecycle_stringent() {
+        // Create a KvManager with 10 blocks capacity
+        let mut manager = KvManager::new(10, 16);
+
+        // Helper function to use multiple blocks
+        fn use_blocks(manager: &mut KvManager, ids: Vec<u64>) {
path: lib/llm/src/mocker/kv_manager.rs
user: login: jthomson04
body: Might be cleaner to have these functions capture the manager instead of passing it as an argument
created_at: 2025-05-19T19:55:46Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,610 @@
+// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use crate::kv_router::protocols::ForwardPassMetrics;
+use crate::mocker::evictor::LRUEvictor;
+use crate::mocker::kv_manager::KvManager;
+use crate::mocker::protocols::DirectRequest;
+use crate::mocker::protocols::{MoveBlock, MoveBlockResponse, PrefillCost, UniqueBlock};
+use crate::mocker::sequence::ActiveSequence;
+use std::collections::HashMap;
+use std::collections::VecDeque;
+use std::sync::Arc;
+use tokio::sync::{mpsc, Mutex};
+use tokio::task::JoinHandle;
+use tokio::time::{interval, Duration};
+use tokio_util::sync::CancellationToken;
+use uuid::Uuid;
+
+/// Enum representing either a direct request or an active sequence
+pub enum Request {
+    Direct(DirectRequest),
+    Active(ActiveSequence),
+}
+struct SchedulerState {
+    waiting: VecDeque<Uuid>,
+    ready: VecDeque<Uuid>,
+    running: LRUEvictor<Uuid>,
+    requests: HashMap<Uuid, Request>,
+    prefill_costs: HashMap<Uuid, Option<PrefillCost>>,
+}
+
+impl Default for SchedulerState {
+    fn default() -> Self {
+        Self {
+            waiting: VecDeque::new(),
+            ready: VecDeque::new(),
+            running: LRUEvictor::new(),
+            requests: HashMap::new(),
+            prefill_costs: HashMap::new(),
+        }
+    }
+}
+
+impl SchedulerState {
+    /// Create a new UUID for a DirectRequest, add it to requests, and push the UUID to waiting.
+    fn receive(&mut self, request: DirectRequest) -> Uuid {
+        // Use the provided UUID if available, otherwise generate a new one
+        let uuid = request.uuid.unwrap_or_else(Uuid::new_v4);
+
+        // Add the request to the map and waiting queue
+        self.requests.insert(uuid, Request::Direct(request));
+        self.waiting.push_back(uuid);
+        uuid
+    }
+
+    /// Get the next UUID from ready or waiting queue and its associated Request.
+    /// Returns from ready if not empty, otherwise from waiting, or None if both are empty.
+    /// Also removes the Request from the requests HashMap.
+    fn next(&mut self) -> Option<(Uuid, Request)> {
+        let uuid = self
+            .ready
+            .pop_front()
+            .or_else(|| self.waiting.pop_front())?;
+        let request = self.requests.remove(&uuid)?;
+        Some((uuid, request))
+    }
+
+    /// Move a UUID and its Request to the ready queue.
+    fn make_ready(&mut self, uuid: Uuid, active_seq: ActiveSequence) {
+        self.requests.insert(uuid, Request::Active(active_seq));
+        self.ready.push_back(uuid);
+    }
+
+    /// Schedule the request with the given UUID.
+    /// Returns the creation signal from the ActiveSequence.
+    fn run(&mut self, uuid: Uuid, active_seq: ActiveSequence) -> MoveBlock {
+        // Insert the request into the map
+        self.requests.insert(uuid, Request::Active(active_seq));
+
+        // Get the creation signal
+        let Some(Request::Active(sequence)) = self.requests.get(&uuid) else {
+            panic!("Failed to get ActiveSequence for UUID");
+        };
+        let Some(signal) = sequence.creation_signal() else {
+            panic!("Failed to get creation signal from ActiveSequence");
+        };
+
+        // Add to running requests
+        self.running.insert(uuid);
+        signal.clone()
+    }
+
+    /// Set the prefill cost for a UUID
+    fn set_prefill_cost(&mut self, uuid: Uuid, cost: Option<PrefillCost>) {
+        self.prefill_costs.insert(uuid, cost);
+    }
+
+    /// Get the prefill compute value for a UUID if available
+    fn get_prefill_compute(&self, uuid: &Uuid) -> Option<f64> {
+        self.prefill_costs
+            .get(uuid)
+            .and_then(|cost| cost.as_ref())
+            .map(|cost| cost.prefill_compute)
+    }
+
+    /// Calculate the current running batched tokens
+    fn num_batched_tokens(&self) -> usize {
+        self.prefill_costs
+            .values()
+            .map(|cost| match cost {
+                Some(cost) => cost.new_tokens,
+                None => 1,
+            })
+            .sum()
+    }
+
+    /// Remove a UUID and its associated Request from collections.
+    fn complete(&mut self, uuid: &Uuid) {
+        self.running.remove(uuid);
+        self.requests.remove(uuid);
+        self.prefill_costs.remove(uuid);
+    }
+
+    /// Preempt the oldest running request by evicting it from running, resetting the sequence,
+    /// and adding it back to the waiting queue.
+    /// Returns the signal from reset_with_signal or None if no requests are running.
+    fn preempt(&mut self) -> Option<Vec<MoveBlock>> {
+        // Evict the oldest UUID from running
+        let uuid = self.running.evict()?;
+        eprintln!("Request {} will be preempted", uuid);
+
+        // Remove the request from the requests HashMap and ensure it's an ActiveSequence
+        let request = self.requests.remove(&uuid)?;
+
+        // Remove the prefill cost to force recomputation
+        self.prefill_costs.remove(&uuid);
+
+        // Extract the ActiveSequence from the Request enum
+        let Request::Active(active_sequence) = request else {
+            panic!("Expected ActiveSequence in running queue")
+        };
+
+        // Reset the sequence and get the new sequence and signal
+        let (new_sequence, signals) = active_sequence.reset_with_signal();
+
+        // Insert the new sequence back into the requests map and add to waiting queue
+        self.requests.insert(uuid, Request::Active(new_sequence));
+        self.waiting.push_back(uuid);
+
+        Some(signals)
+    }
+}
+
+/// Manages scheduling of requests using KvManager resources
+pub struct Scheduler {
+    state: Arc<Mutex<SchedulerState>>,
+    kv_manager: Arc<Mutex<KvManager>>,
+    token_capacity: usize,
+    watermark: f64,
+    block_size: usize,
+    chunk_size: usize,
+    background_handle: Option<JoinHandle<()>>,
+    request_tx: mpsc::Sender<DirectRequest>,
+    cancellation_token: CancellationToken,
+}
+
+impl Scheduler {
+    /// Create a new Scheduler with the given parameters
+    pub fn new(
+        kv_capacity: usize,
+        watermark: f64,
+        block_size: usize,
+        chunk_size: Option<usize>,
+        output_tx: Option<mpsc::Sender<Uuid>>,
+        cancellation_token: Option<CancellationToken>,
+    ) -> Self {
+        // Create KvManager internally
+        let kv_manager = KvManager::new(kv_capacity, block_size);
+
+        let token_capacity: usize = 8192;
+        let state = Arc::new(Mutex::new(SchedulerState::default()));
path: lib/llm/src/mocker/scheduler.rs
user: login: jthomson04
body: Is there any way we can split `state` into multiple things? Seems like a ton of functionality all requires these  locks. At the very least, might be worthwhile to have a `tokio::sync::RwLock`
created_at: 2025-05-19T20:04:32Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,610 @@
+// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use crate::kv_router::protocols::ForwardPassMetrics;
+use crate::mocker::evictor::LRUEvictor;
+use crate::mocker::kv_manager::KvManager;
+use crate::mocker::protocols::DirectRequest;
+use crate::mocker::protocols::{MoveBlock, MoveBlockResponse, PrefillCost, UniqueBlock};
+use crate::mocker::sequence::ActiveSequence;
+use std::collections::HashMap;
+use std::collections::VecDeque;
+use std::sync::Arc;
+use tokio::sync::{mpsc, Mutex};
+use tokio::task::JoinHandle;
+use tokio::time::{interval, Duration};
+use tokio_util::sync::CancellationToken;
+use uuid::Uuid;
+
+/// Enum representing either a direct request or an active sequence
+pub enum Request {
+    Direct(DirectRequest),
+    Active(ActiveSequence),
+}
+struct SchedulerState {
+    waiting: VecDeque<Uuid>,
+    ready: VecDeque<Uuid>,
+    running: LRUEvictor<Uuid>,
+    requests: HashMap<Uuid, Request>,
+    prefill_costs: HashMap<Uuid, Option<PrefillCost>>,
+}
+
+impl Default for SchedulerState {
+    fn default() -> Self {
+        Self {
+            waiting: VecDeque::new(),
+            ready: VecDeque::new(),
+            running: LRUEvictor::new(),
+            requests: HashMap::new(),
+            prefill_costs: HashMap::new(),
+        }
+    }
+}
+
+impl SchedulerState {
+    /// Create a new UUID for a DirectRequest, add it to requests, and push the UUID to waiting.
+    fn receive(&mut self, request: DirectRequest) -> Uuid {
+        // Use the provided UUID if available, otherwise generate a new one
+        let uuid = request.uuid.unwrap_or_else(Uuid::new_v4);
+
+        // Add the request to the map and waiting queue
+        self.requests.insert(uuid, Request::Direct(request));
+        self.waiting.push_back(uuid);
+        uuid
+    }
+
+    /// Get the next UUID from ready or waiting queue and its associated Request.
+    /// Returns from ready if not empty, otherwise from waiting, or None if both are empty.
+    /// Also removes the Request from the requests HashMap.
+    fn next(&mut self) -> Option<(Uuid, Request)> {
+        let uuid = self
+            .ready
+            .pop_front()
+            .or_else(|| self.waiting.pop_front())?;
+        let request = self.requests.remove(&uuid)?;
+        Some((uuid, request))
+    }
+
+    /// Move a UUID and its Request to the ready queue.
+    fn make_ready(&mut self, uuid: Uuid, active_seq: ActiveSequence) {
+        self.requests.insert(uuid, Request::Active(active_seq));
+        self.ready.push_back(uuid);
+    }
+
+    /// Schedule the request with the given UUID.
+    /// Returns the creation signal from the ActiveSequence.
+    fn run(&mut self, uuid: Uuid, active_seq: ActiveSequence) -> MoveBlock {
+        // Insert the request into the map
+        self.requests.insert(uuid, Request::Active(active_seq));
+
+        // Get the creation signal
+        let Some(Request::Active(sequence)) = self.requests.get(&uuid) else {
+            panic!("Failed to get ActiveSequence for UUID");
+        };
+        let Some(signal) = sequence.creation_signal() else {
+            panic!("Failed to get creation signal from ActiveSequence");
+        };
+
+        // Add to running requests
+        self.running.insert(uuid);
+        signal.clone()
+    }
+
+    /// Set the prefill cost for a UUID
+    fn set_prefill_cost(&mut self, uuid: Uuid, cost: Option<PrefillCost>) {
+        self.prefill_costs.insert(uuid, cost);
+    }
+
+    /// Get the prefill compute value for a UUID if available
+    fn get_prefill_compute(&self, uuid: &Uuid) -> Option<f64> {
+        self.prefill_costs
+            .get(uuid)
+            .and_then(|cost| cost.as_ref())
+            .map(|cost| cost.prefill_compute)
+    }
+
+    /// Calculate the current running batched tokens
+    fn num_batched_tokens(&self) -> usize {
+        self.prefill_costs
+            .values()
+            .map(|cost| match cost {
+                Some(cost) => cost.new_tokens,
+                None => 1,
+            })
+            .sum()
+    }
+
+    /// Remove a UUID and its associated Request from collections.
+    fn complete(&mut self, uuid: &Uuid) {
+        self.running.remove(uuid);
+        self.requests.remove(uuid);
+        self.prefill_costs.remove(uuid);
+    }
+
+    /// Preempt the oldest running request by evicting it from running, resetting the sequence,
+    /// and adding it back to the waiting queue.
+    /// Returns the signal from reset_with_signal or None if no requests are running.
+    fn preempt(&mut self) -> Option<Vec<MoveBlock>> {
+        // Evict the oldest UUID from running
+        let uuid = self.running.evict()?;
+        eprintln!("Request {} will be preempted", uuid);
+
+        // Remove the request from the requests HashMap and ensure it's an ActiveSequence
+        let request = self.requests.remove(&uuid)?;
+
+        // Remove the prefill cost to force recomputation
+        self.prefill_costs.remove(&uuid);
+
+        // Extract the ActiveSequence from the Request enum
+        let Request::Active(active_sequence) = request else {
+            panic!("Expected ActiveSequence in running queue")
+        };
+
+        // Reset the sequence and get the new sequence and signal
+        let (new_sequence, signals) = active_sequence.reset_with_signal();
+
+        // Insert the new sequence back into the requests map and add to waiting queue
+        self.requests.insert(uuid, Request::Active(new_sequence));
+        self.waiting.push_back(uuid);
+
+        Some(signals)
+    }
+}
+
+/// Manages scheduling of requests using KvManager resources
+pub struct Scheduler {
+    state: Arc<Mutex<SchedulerState>>,
+    kv_manager: Arc<Mutex<KvManager>>,
+    token_capacity: usize,
+    watermark: f64,
+    block_size: usize,
+    chunk_size: usize,
+    background_handle: Option<JoinHandle<()>>,
+    request_tx: mpsc::Sender<DirectRequest>,
+    cancellation_token: CancellationToken,
+}
+
+impl Scheduler {
+    /// Create a new Scheduler with the given parameters
+    pub fn new(
+        kv_capacity: usize,
+        watermark: f64,
+        block_size: usize,
+        chunk_size: Option<usize>,
+        output_tx: Option<mpsc::Sender<Uuid>>,
+        cancellation_token: Option<CancellationToken>,
+    ) -> Self {
+        // Create KvManager internally
+        let kv_manager = KvManager::new(kv_capacity, block_size);
+
+        let token_capacity: usize = 8192;
+        let state = Arc::new(Mutex::new(SchedulerState::default()));
+
+        let kv_manager = Arc::new(Mutex::new(kv_manager));
+        let chunk_size = chunk_size.unwrap_or(256);
+
+        // Create channel for request handling
+        let (request_tx, mut request_rx) = mpsc::channel::<DirectRequest>(1024);
+
+        // Use provided cancellation token or create new one
+        let cancellation_token = cancellation_token.unwrap_or_default();
+        let token_clone = cancellation_token.clone();
+
+        // Create a clone for the background task
+        let state_clone = state.clone();
+        let kv_manager_clone = kv_manager.clone();
+        let output_tx_clone = output_tx.clone();
+
+        // Spawn main background task with cancellation token
+        let background_handle = tokio::spawn(async move {
path: lib/llm/src/mocker/scheduler.rs
user: login: jthomson04
body: This will fail if not called within a tokio runtime. Maybe check if you're inside one already, and create one if not.
created_at: 2025-05-19T20:05:21Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,610 @@
+// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use crate::kv_router::protocols::ForwardPassMetrics;
+use crate::mocker::evictor::LRUEvictor;
+use crate::mocker::kv_manager::KvManager;
+use crate::mocker::protocols::DirectRequest;
+use crate::mocker::protocols::{MoveBlock, MoveBlockResponse, PrefillCost, UniqueBlock};
+use crate::mocker::sequence::ActiveSequence;
+use std::collections::HashMap;
+use std::collections::VecDeque;
+use std::sync::Arc;
+use tokio::sync::{mpsc, Mutex};
+use tokio::task::JoinHandle;
+use tokio::time::{interval, Duration};
+use tokio_util::sync::CancellationToken;
+use uuid::Uuid;
+
+/// Enum representing either a direct request or an active sequence
+pub enum Request {
+    Direct(DirectRequest),
+    Active(ActiveSequence),
+}
+struct SchedulerState {
+    waiting: VecDeque<Uuid>,
+    ready: VecDeque<Uuid>,
+    running: LRUEvictor<Uuid>,
+    requests: HashMap<Uuid, Request>,
+    prefill_costs: HashMap<Uuid, Option<PrefillCost>>,
+}
+
+impl Default for SchedulerState {
+    fn default() -> Self {
+        Self {
+            waiting: VecDeque::new(),
+            ready: VecDeque::new(),
+            running: LRUEvictor::new(),
+            requests: HashMap::new(),
+            prefill_costs: HashMap::new(),
+        }
+    }
+}
+
+impl SchedulerState {
+    /// Create a new UUID for a DirectRequest, add it to requests, and push the UUID to waiting.
+    fn receive(&mut self, request: DirectRequest) -> Uuid {
+        // Use the provided UUID if available, otherwise generate a new one
+        let uuid = request.uuid.unwrap_or_else(Uuid::new_v4);
+
+        // Add the request to the map and waiting queue
+        self.requests.insert(uuid, Request::Direct(request));
+        self.waiting.push_back(uuid);
+        uuid
+    }
+
+    /// Get the next UUID from ready or waiting queue and its associated Request.
+    /// Returns from ready if not empty, otherwise from waiting, or None if both are empty.
+    /// Also removes the Request from the requests HashMap.
+    fn next(&mut self) -> Option<(Uuid, Request)> {
+        let uuid = self
+            .ready
+            .pop_front()
+            .or_else(|| self.waiting.pop_front())?;
+        let request = self.requests.remove(&uuid)?;
+        Some((uuid, request))
+    }
+
+    /// Move a UUID and its Request to the ready queue.
+    fn make_ready(&mut self, uuid: Uuid, active_seq: ActiveSequence) {
+        self.requests.insert(uuid, Request::Active(active_seq));
+        self.ready.push_back(uuid);
+    }
+
+    /// Schedule the request with the given UUID.
+    /// Returns the creation signal from the ActiveSequence.
+    fn run(&mut self, uuid: Uuid, active_seq: ActiveSequence) -> MoveBlock {
+        // Insert the request into the map
+        self.requests.insert(uuid, Request::Active(active_seq));
+
+        // Get the creation signal
+        let Some(Request::Active(sequence)) = self.requests.get(&uuid) else {
+            panic!("Failed to get ActiveSequence for UUID");
+        };
+        let Some(signal) = sequence.creation_signal() else {
+            panic!("Failed to get creation signal from ActiveSequence");
+        };
+
+        // Add to running requests
+        self.running.insert(uuid);
+        signal.clone()
+    }
+
+    /// Set the prefill cost for a UUID
+    fn set_prefill_cost(&mut self, uuid: Uuid, cost: Option<PrefillCost>) {
+        self.prefill_costs.insert(uuid, cost);
+    }
+
+    /// Get the prefill compute value for a UUID if available
+    fn get_prefill_compute(&self, uuid: &Uuid) -> Option<f64> {
+        self.prefill_costs
+            .get(uuid)
+            .and_then(|cost| cost.as_ref())
+            .map(|cost| cost.prefill_compute)
+    }
+
+    /// Calculate the current running batched tokens
+    fn num_batched_tokens(&self) -> usize {
+        self.prefill_costs
+            .values()
+            .map(|cost| match cost {
+                Some(cost) => cost.new_tokens,
+                None => 1,
+            })
+            .sum()
+    }
+
+    /// Remove a UUID and its associated Request from collections.
+    fn complete(&mut self, uuid: &Uuid) {
+        self.running.remove(uuid);
+        self.requests.remove(uuid);
+        self.prefill_costs.remove(uuid);
+    }
+
+    /// Preempt the oldest running request by evicting it from running, resetting the sequence,
+    /// and adding it back to the waiting queue.
+    /// Returns the signal from reset_with_signal or None if no requests are running.
+    fn preempt(&mut self) -> Option<Vec<MoveBlock>> {
+        // Evict the oldest UUID from running
+        let uuid = self.running.evict()?;
+        eprintln!("Request {} will be preempted", uuid);
+
+        // Remove the request from the requests HashMap and ensure it's an ActiveSequence
+        let request = self.requests.remove(&uuid)?;
+
+        // Remove the prefill cost to force recomputation
+        self.prefill_costs.remove(&uuid);
+
+        // Extract the ActiveSequence from the Request enum
+        let Request::Active(active_sequence) = request else {
+            panic!("Expected ActiveSequence in running queue")
+        };
+
+        // Reset the sequence and get the new sequence and signal
+        let (new_sequence, signals) = active_sequence.reset_with_signal();
+
+        // Insert the new sequence back into the requests map and add to waiting queue
+        self.requests.insert(uuid, Request::Active(new_sequence));
+        self.waiting.push_back(uuid);
+
+        Some(signals)
+    }
+}
+
+/// Manages scheduling of requests using KvManager resources
+pub struct Scheduler {
+    state: Arc<Mutex<SchedulerState>>,
+    kv_manager: Arc<Mutex<KvManager>>,
+    token_capacity: usize,
+    watermark: f64,
+    block_size: usize,
+    chunk_size: usize,
+    background_handle: Option<JoinHandle<()>>,
+    request_tx: mpsc::Sender<DirectRequest>,
+    cancellation_token: CancellationToken,
+}
+
+impl Scheduler {
+    /// Create a new Scheduler with the given parameters
+    pub fn new(
+        kv_capacity: usize,
+        watermark: f64,
+        block_size: usize,
+        chunk_size: Option<usize>,
+        output_tx: Option<mpsc::Sender<Uuid>>,
+        cancellation_token: Option<CancellationToken>,
+    ) -> Self {
+        // Create KvManager internally
+        let kv_manager = KvManager::new(kv_capacity, block_size);
+
+        let token_capacity: usize = 8192;
+        let state = Arc::new(Mutex::new(SchedulerState::default()));
+
+        let kv_manager = Arc::new(Mutex::new(kv_manager));
+        let chunk_size = chunk_size.unwrap_or(256);
+
+        // Create channel for request handling
+        let (request_tx, mut request_rx) = mpsc::channel::<DirectRequest>(1024);
+
+        // Use provided cancellation token or create new one
+        let cancellation_token = cancellation_token.unwrap_or_default();
+        let token_clone = cancellation_token.clone();
+
+        // Create a clone for the background task
+        let state_clone = state.clone();
+        let kv_manager_clone = kv_manager.clone();
+        let output_tx_clone = output_tx.clone();
+
+        // Spawn main background task with cancellation token
+        let background_handle = tokio::spawn(async move {
+            let mut schedule_interval = interval(Duration::from_millis(5));
+            let mut simulate_interval = interval(Duration::from_millis(1));
+
+            loop {
+                tokio::select! {
+                    biased;
+
+                    // Enqueue new request
+                    Some(request) = request_rx.recv() => {
+                        let mut state = state_clone.lock().await;
+                        state.receive(request);
+                    }
+
+                    // Try Scheduling Requests
+                    _ = schedule_interval.tick() => {
+                        let mut state_guard = state_clone.lock().await;
+                        let mut kv_manager_guard = kv_manager_clone.lock().await;
+
+                        // Process DirectRequests, converting them to ActiveSequence and scheduling them until we can't
+                        // schedule anymore.
+                        while let Some((uuid, request)) = state_guard.next() {
+                            let active_sequence = get_active_sequence(request, block_size, chunk_size);
+
+                            // Calculate token budget using new_tokens from PrefillCost
+                            let total_prefill_tokens = state_guard.num_batched_tokens();
+                            let tokens_budget = token_capacity.saturating_sub(total_prefill_tokens);
+
+                            // Check if it can be scheduled
+                            let Some(prefill_cost) = kv_manager_guard.try_schedule(&active_sequence, watermark, tokens_budget) else {
+                                state_guard.make_ready(uuid, active_sequence);
+                                break;
+                            };
+
+                            // Get creation signal and schedule the request
+                            let signal = state_guard.run(uuid, active_sequence);
+                            kv_manager_guard.process(&signal);
+                            state_guard.set_prefill_cost(uuid, Some(prefill_cost));
+                        }
+                    }
+
+                    // Check for cancellation
+                    _ = token_clone.cancelled() => {
+                        break;
+                    }
+
+                    // Simulate running requests (prefill + decode)
+                    _ = simulate_interval.tick() => {
+                        let mut state_guard = state_clone.lock().await;
+                        let mut kv_manager_guard = kv_manager_clone.lock().await;
+
+                        // Process each running request
+                        let mut total_sleep_duration = Duration::from_millis(0);
+                        let uuids: Vec<Uuid> = state_guard.running.keys().cloned().collect();
+                        for uuid in uuids {
+                            // Check if UUID is still in running_requests, if not skip this iteration
+                            if !state_guard.running.contains(&uuid) {
+                                continue;
+                            }
+
+                            // Get prefill compute value first
+                            let prefill_compute = state_guard.get_prefill_compute(&uuid);
+
+                            // Get the active sequence for this UUID
+                            let sequence = state_guard.requests.get_mut(&uuid)
+                                .and_then(|req| if let Request::Active(seq) = req { Some(seq) } else { None })
+                                .expect("UUID in running_requests must have a corresponding active sequence");
+
+                            // Generate token and get signals
+                            let signals = sequence.generate();
+
+                            // Accumulate sleep duration based on prefill_compute if available
+                            let sleep_ms = if let Some(compute) = prefill_compute {
+                                (compute / 131072.0) as u64  // TODO: A magic scaling to model prefill compute
+                            } else {
+                                1  // FIXME: A dummy number to simulate ITL. Need some work.
+                            };
+                            total_sleep_duration += Duration::from_millis(sleep_ms);
+
+                            // Process all signals with the KvManager
+                            if process_signals(&mut kv_manager_guard, &signals) == MoveBlockResponse::Failure {
+                                sequence.partial_reset();
+
+                                // free_signal derefs the preempted blocks
+                                let Some(free_signal) = state_guard.preempt() else {
+                                    panic!("Failed to acquire signal to free KV blocks from preemption");
+                                };
+
+                                for signal in free_signal {
+                                    kv_manager_guard.process(&signal);
+                                }
+                                continue;
+                            }
+
+                            // Send UUID notification for each generated token
+                            // TODO: hook this up to an AsyncEngine
+                            if let Some(tx) = &output_tx_clone {
+                                let _ = tx.try_send(uuid);
+                            }
+
+                            // Check if we're done after generating
+                            if sequence.generated_tokens() >= sequence.max_output_tokens() {
+                                state_guard.complete(&uuid);
+                                continue;
+                            }
+
+                            // Transition to decode (no prefill cost)
+                            if sequence.generated_tokens() == 1 {
+                                state_guard.set_prefill_cost(uuid, None);
+                            }
+                        }
+
+                        // Sleep once for the accumulated duration
+                        if total_sleep_duration.as_millis() > 0 {
+                            tokio::time::sleep(total_sleep_duration).await;
+                        }
+                    }
+                }
+            }
+        });
+
+        Self {
+            state,
+            kv_manager,
+            token_capacity,
+            watermark,
+            block_size,
+            chunk_size,
+            background_handle: Some(background_handle),
+            request_tx,
+            cancellation_token,
+        }
+    }
+
+    /// Add a new request to the waiting queue
+    pub async fn receive(&self, request: DirectRequest) {
+        let _ = self.request_tx.send(request).await;
+    }
+
+    /// Get the count of waiting requests
+    pub async fn waiting_count(&self) -> usize {
+        let state = self.state.lock().await;
+        state.waiting.len()
+    }
+
+    /// Get the count of running requests
+    pub async fn running_count(&self) -> usize {
+        let state = self.state.lock().await;
+        state.running.len()
+    }
+
+    /// Get the current capacity of the KvManager
+    pub async fn kv_usage_perc(&self) -> f64 {
+        let kv_manager = self.kv_manager.lock().await;
+        kv_manager.current_capacity_perc()
+    }
+
+    /// Returns forward pass metrics for monitoring purposes
+    pub async fn get_forward_pass_metrics(&self) -> ForwardPassMetrics {
+        let state = self.state.lock().await;
+        let kv_manager = self.kv_manager.lock().await;
+
+        // Get the active blocks and total capacity from KvManager
+        let active_blocks_count = kv_manager.active_blocks().len() as u64;
+        let total_capacity = kv_manager.max_capacity() as u64;
+
+        // Calculate GPU cache usage percentage
+        let gpu_cache_usage_perc = if total_capacity > 0 {
+            active_blocks_count as f32 / total_capacity as f32
+        } else {
+            0.0
+        };
+
+        ForwardPassMetrics {
+            request_active_slots: state.running.len() as u64,
+            request_total_slots: 420, // Dummy value as specified
+            kv_active_blocks: active_blocks_count,
+            kv_total_blocks: total_capacity,
+            num_requests_waiting: state.waiting.len() as u64,
+            gpu_cache_usage_perc,
+            gpu_prefix_cache_hit_rate: 0.0, // Placeholder value as specified
+        }
+    }
+}
+
+// Implement Clone for Scheduler to support sharing between tasks
+impl Clone for Scheduler {
+    fn clone(&self) -> Self {
+        Self {
+            state: self.state.clone(),
+            kv_manager: self.kv_manager.clone(),
+            token_capacity: self.token_capacity,
+            watermark: self.watermark,
+            block_size: self.block_size,
+            chunk_size: self.chunk_size,
+            background_handle: None,
path: lib/llm/src/mocker/scheduler.rs
user: login: jthomson04
body: This seems problematic. What if we clone the `Scheduler`, then drop the original instance? Then the second cloned instance will get messed up.
created_at: 2025-05-19T20:09:52Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,610 @@
+// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use crate::kv_router::protocols::ForwardPassMetrics;
+use crate::mocker::evictor::LRUEvictor;
+use crate::mocker::kv_manager::KvManager;
+use crate::mocker::protocols::DirectRequest;
+use crate::mocker::protocols::{MoveBlock, MoveBlockResponse, PrefillCost, UniqueBlock};
+use crate::mocker::sequence::ActiveSequence;
+use std::collections::HashMap;
+use std::collections::VecDeque;
+use std::sync::Arc;
+use tokio::sync::{mpsc, Mutex};
+use tokio::task::JoinHandle;
+use tokio::time::{interval, Duration};
+use tokio_util::sync::CancellationToken;
+use uuid::Uuid;
+
+/// Enum representing either a direct request or an active sequence
+pub enum Request {
+    Direct(DirectRequest),
+    Active(ActiveSequence),
+}
+struct SchedulerState {
+    waiting: VecDeque<Uuid>,
+    ready: VecDeque<Uuid>,
+    running: LRUEvictor<Uuid>,
+    requests: HashMap<Uuid, Request>,
+    prefill_costs: HashMap<Uuid, Option<PrefillCost>>,
+}
+
+impl Default for SchedulerState {
+    fn default() -> Self {
+        Self {
+            waiting: VecDeque::new(),
+            ready: VecDeque::new(),
+            running: LRUEvictor::new(),
+            requests: HashMap::new(),
+            prefill_costs: HashMap::new(),
+        }
+    }
+}
+
+impl SchedulerState {
+    /// Create a new UUID for a DirectRequest, add it to requests, and push the UUID to waiting.
+    fn receive(&mut self, request: DirectRequest) -> Uuid {
+        // Use the provided UUID if available, otherwise generate a new one
+        let uuid = request.uuid.unwrap_or_else(Uuid::new_v4);
+
+        // Add the request to the map and waiting queue
+        self.requests.insert(uuid, Request::Direct(request));
+        self.waiting.push_back(uuid);
+        uuid
+    }
+
+    /// Get the next UUID from ready or waiting queue and its associated Request.
+    /// Returns from ready if not empty, otherwise from waiting, or None if both are empty.
+    /// Also removes the Request from the requests HashMap.
+    fn next(&mut self) -> Option<(Uuid, Request)> {
+        let uuid = self
+            .ready
+            .pop_front()
+            .or_else(|| self.waiting.pop_front())?;
+        let request = self.requests.remove(&uuid)?;
+        Some((uuid, request))
+    }
+
+    /// Move a UUID and its Request to the ready queue.
+    fn make_ready(&mut self, uuid: Uuid, active_seq: ActiveSequence) {
+        self.requests.insert(uuid, Request::Active(active_seq));
+        self.ready.push_back(uuid);
+    }
+
+    /// Schedule the request with the given UUID.
+    /// Returns the creation signal from the ActiveSequence.
+    fn run(&mut self, uuid: Uuid, active_seq: ActiveSequence) -> MoveBlock {
+        // Insert the request into the map
+        self.requests.insert(uuid, Request::Active(active_seq));
+
+        // Get the creation signal
+        let Some(Request::Active(sequence)) = self.requests.get(&uuid) else {
+            panic!("Failed to get ActiveSequence for UUID");
+        };
+        let Some(signal) = sequence.creation_signal() else {
+            panic!("Failed to get creation signal from ActiveSequence");
+        };
+
+        // Add to running requests
+        self.running.insert(uuid);
+        signal.clone()
+    }
+
+    /// Set the prefill cost for a UUID
+    fn set_prefill_cost(&mut self, uuid: Uuid, cost: Option<PrefillCost>) {
+        self.prefill_costs.insert(uuid, cost);
+    }
+
+    /// Get the prefill compute value for a UUID if available
+    fn get_prefill_compute(&self, uuid: &Uuid) -> Option<f64> {
+        self.prefill_costs
+            .get(uuid)
+            .and_then(|cost| cost.as_ref())
+            .map(|cost| cost.prefill_compute)
+    }
+
+    /// Calculate the current running batched tokens
+    fn num_batched_tokens(&self) -> usize {
+        self.prefill_costs
+            .values()
+            .map(|cost| match cost {
+                Some(cost) => cost.new_tokens,
+                None => 1,
+            })
+            .sum()
+    }
+
+    /// Remove a UUID and its associated Request from collections.
+    fn complete(&mut self, uuid: &Uuid) {
+        self.running.remove(uuid);
+        self.requests.remove(uuid);
+        self.prefill_costs.remove(uuid);
+    }
+
+    /// Preempt the oldest running request by evicting it from running, resetting the sequence,
+    /// and adding it back to the waiting queue.
+    /// Returns the signal from reset_with_signal or None if no requests are running.
+    fn preempt(&mut self) -> Option<Vec<MoveBlock>> {
+        // Evict the oldest UUID from running
+        let uuid = self.running.evict()?;
+        eprintln!("Request {} will be preempted", uuid);
+
+        // Remove the request from the requests HashMap and ensure it's an ActiveSequence
+        let request = self.requests.remove(&uuid)?;
+
+        // Remove the prefill cost to force recomputation
+        self.prefill_costs.remove(&uuid);
+
+        // Extract the ActiveSequence from the Request enum
+        let Request::Active(active_sequence) = request else {
+            panic!("Expected ActiveSequence in running queue")
+        };
+
+        // Reset the sequence and get the new sequence and signal
+        let (new_sequence, signals) = active_sequence.reset_with_signal();
+
+        // Insert the new sequence back into the requests map and add to waiting queue
+        self.requests.insert(uuid, Request::Active(new_sequence));
+        self.waiting.push_back(uuid);
+
+        Some(signals)
+    }
+}
+
+/// Manages scheduling of requests using KvManager resources
+pub struct Scheduler {
+    state: Arc<Mutex<SchedulerState>>,
+    kv_manager: Arc<Mutex<KvManager>>,
+    token_capacity: usize,
+    watermark: f64,
+    block_size: usize,
+    chunk_size: usize,
+    background_handle: Option<JoinHandle<()>>,
+    request_tx: mpsc::Sender<DirectRequest>,
+    cancellation_token: CancellationToken,
+}
+
+impl Scheduler {
+    /// Create a new Scheduler with the given parameters
+    pub fn new(
+        kv_capacity: usize,
+        watermark: f64,
+        block_size: usize,
+        chunk_size: Option<usize>,
+        output_tx: Option<mpsc::Sender<Uuid>>,
+        cancellation_token: Option<CancellationToken>,
+    ) -> Self {
+        // Create KvManager internally
+        let kv_manager = KvManager::new(kv_capacity, block_size);
+
+        let token_capacity: usize = 8192;
+        let state = Arc::new(Mutex::new(SchedulerState::default()));
+
+        let kv_manager = Arc::new(Mutex::new(kv_manager));
+        let chunk_size = chunk_size.unwrap_or(256);
+
+        // Create channel for request handling
+        let (request_tx, mut request_rx) = mpsc::channel::<DirectRequest>(1024);
+
+        // Use provided cancellation token or create new one
+        let cancellation_token = cancellation_token.unwrap_or_default();
+        let token_clone = cancellation_token.clone();
+
+        // Create a clone for the background task
+        let state_clone = state.clone();
+        let kv_manager_clone = kv_manager.clone();
+        let output_tx_clone = output_tx.clone();
+
+        // Spawn main background task with cancellation token
+        let background_handle = tokio::spawn(async move {
+            let mut schedule_interval = interval(Duration::from_millis(5));
+            let mut simulate_interval = interval(Duration::from_millis(1));
+
+            loop {
+                tokio::select! {
+                    biased;
+
+                    // Enqueue new request
+                    Some(request) = request_rx.recv() => {
+                        let mut state = state_clone.lock().await;
+                        state.receive(request);
+                    }
+
+                    // Try Scheduling Requests
+                    _ = schedule_interval.tick() => {
+                        let mut state_guard = state_clone.lock().await;
+                        let mut kv_manager_guard = kv_manager_clone.lock().await;
+
+                        // Process DirectRequests, converting them to ActiveSequence and scheduling them until we can't
+                        // schedule anymore.
+                        while let Some((uuid, request)) = state_guard.next() {
+                            let active_sequence = get_active_sequence(request, block_size, chunk_size);
+
+                            // Calculate token budget using new_tokens from PrefillCost
+                            let total_prefill_tokens = state_guard.num_batched_tokens();
+                            let tokens_budget = token_capacity.saturating_sub(total_prefill_tokens);
+
+                            // Check if it can be scheduled
+                            let Some(prefill_cost) = kv_manager_guard.try_schedule(&active_sequence, watermark, tokens_budget) else {
+                                state_guard.make_ready(uuid, active_sequence);
+                                break;
+                            };
+
+                            // Get creation signal and schedule the request
+                            let signal = state_guard.run(uuid, active_sequence);
+                            kv_manager_guard.process(&signal);
+                            state_guard.set_prefill_cost(uuid, Some(prefill_cost));
+                        }
+                    }
+
+                    // Check for cancellation
+                    _ = token_clone.cancelled() => {
+                        break;
+                    }
+
+                    // Simulate running requests (prefill + decode)
+                    _ = simulate_interval.tick() => {
+                        let mut state_guard = state_clone.lock().await;
+                        let mut kv_manager_guard = kv_manager_clone.lock().await;
+
+                        // Process each running request
+                        let mut total_sleep_duration = Duration::from_millis(0);
+                        let uuids: Vec<Uuid> = state_guard.running.keys().cloned().collect();
+                        for uuid in uuids {
+                            // Check if UUID is still in running_requests, if not skip this iteration
+                            if !state_guard.running.contains(&uuid) {
+                                continue;
+                            }
+
+                            // Get prefill compute value first
+                            let prefill_compute = state_guard.get_prefill_compute(&uuid);
+
+                            // Get the active sequence for this UUID
+                            let sequence = state_guard.requests.get_mut(&uuid)
+                                .and_then(|req| if let Request::Active(seq) = req { Some(seq) } else { None })
+                                .expect("UUID in running_requests must have a corresponding active sequence");
+
+                            // Generate token and get signals
+                            let signals = sequence.generate();
+
+                            // Accumulate sleep duration based on prefill_compute if available
+                            let sleep_ms = if let Some(compute) = prefill_compute {
+                                (compute / 131072.0) as u64  // TODO: A magic scaling to model prefill compute
+                            } else {
+                                1  // FIXME: A dummy number to simulate ITL. Need some work.
+                            };
+                            total_sleep_duration += Duration::from_millis(sleep_ms);
+
+                            // Process all signals with the KvManager
+                            if process_signals(&mut kv_manager_guard, &signals) == MoveBlockResponse::Failure {
+                                sequence.partial_reset();
+
+                                // free_signal derefs the preempted blocks
+                                let Some(free_signal) = state_guard.preempt() else {
+                                    panic!("Failed to acquire signal to free KV blocks from preemption");
+                                };
+
+                                for signal in free_signal {
+                                    kv_manager_guard.process(&signal);
+                                }
+                                continue;
+                            }
+
+                            // Send UUID notification for each generated token
+                            // TODO: hook this up to an AsyncEngine
+                            if let Some(tx) = &output_tx_clone {
+                                let _ = tx.try_send(uuid);
+                            }
+
+                            // Check if we're done after generating
+                            if sequence.generated_tokens() >= sequence.max_output_tokens() {
+                                state_guard.complete(&uuid);
+                                continue;
+                            }
+
+                            // Transition to decode (no prefill cost)
+                            if sequence.generated_tokens() == 1 {
+                                state_guard.set_prefill_cost(uuid, None);
+                            }
+                        }
+
+                        // Sleep once for the accumulated duration
+                        if total_sleep_duration.as_millis() > 0 {
+                            tokio::time::sleep(total_sleep_duration).await;
+                        }
+                    }
+                }
+            }
+        });
+
+        Self {
+            state,
+            kv_manager,
+            token_capacity,
+            watermark,
+            block_size,
+            chunk_size,
+            background_handle: Some(background_handle),
+            request_tx,
+            cancellation_token,
+        }
+    }
+
+    /// Add a new request to the waiting queue
+    pub async fn receive(&self, request: DirectRequest) {
+        let _ = self.request_tx.send(request).await;
+    }
+
+    /// Get the count of waiting requests
+    pub async fn waiting_count(&self) -> usize {
+        let state = self.state.lock().await;
+        state.waiting.len()
+    }
+
+    /// Get the count of running requests
+    pub async fn running_count(&self) -> usize {
+        let state = self.state.lock().await;
+        state.running.len()
+    }
+
+    /// Get the current capacity of the KvManager
+    pub async fn kv_usage_perc(&self) -> f64 {
+        let kv_manager = self.kv_manager.lock().await;
+        kv_manager.current_capacity_perc()
+    }
+
+    /// Returns forward pass metrics for monitoring purposes
+    pub async fn get_forward_pass_metrics(&self) -> ForwardPassMetrics {
+        let state = self.state.lock().await;
+        let kv_manager = self.kv_manager.lock().await;
+
+        // Get the active blocks and total capacity from KvManager
+        let active_blocks_count = kv_manager.active_blocks().len() as u64;
+        let total_capacity = kv_manager.max_capacity() as u64;
+
+        // Calculate GPU cache usage percentage
+        let gpu_cache_usage_perc = if total_capacity > 0 {
+            active_blocks_count as f32 / total_capacity as f32
+        } else {
+            0.0
+        };
+
+        ForwardPassMetrics {
+            request_active_slots: state.running.len() as u64,
+            request_total_slots: 420, // Dummy value as specified
+            kv_active_blocks: active_blocks_count,
+            kv_total_blocks: total_capacity,
+            num_requests_waiting: state.waiting.len() as u64,
+            gpu_cache_usage_perc,
+            gpu_prefix_cache_hit_rate: 0.0, // Placeholder value as specified
+        }
+    }
+}
+
+// Implement Clone for Scheduler to support sharing between tasks
+impl Clone for Scheduler {
+    fn clone(&self) -> Self {
+        Self {
+            state: self.state.clone(),
+            kv_manager: self.kv_manager.clone(),
+            token_capacity: self.token_capacity,
+            watermark: self.watermark,
+            block_size: self.block_size,
+            chunk_size: self.chunk_size,
+            background_handle: None,
+            request_tx: self.request_tx.clone(),
+            cancellation_token: self.cancellation_token.clone(),
+        }
+    }
+}
+
+impl Drop for Scheduler {
+    fn drop(&mut self) {
+        // Only the original instance should cancel
+        if self.background_handle.is_some() {
+            self.cancellation_token.cancel();
+        }
+
+        // Abort the background task if this is the original instance
+        if let Some(handle) = self.background_handle.take() {
+            handle.abort();
+        }
+    }
+}
+
+/// Convert a Request to an ActiveSequence
+fn get_active_sequence(request: Request, block_size: usize, chunk_size: usize) -> ActiveSequence {
+    if let Request::Active(active_seq) = request {
+        return active_seq;
+    }
+
+    let Request::Direct(direct_request) = request else {
+        unreachable!("Request must be either Direct or Active");
+    };
+
+    ActiveSequence::new(
+        direct_request.tokens,
+        direct_request.max_output_tokens,
+        Some(block_size),
+        Some(chunk_size),
+    )
+}
+
+/// Processes MoveBlock signals with the KvManager.
+///
+/// When a signal fails, this function verifies that the failure is for an expected case:
+/// specifically a single signal attempting to create a single partial (generation) block.
+/// This validation is important because in normal operation, the only legitimate failure
+/// case should be when trying to acquire a new generation block - any other failures would
+/// indicate an unexpected state in the system.
+fn process_signals(
+    kv_manager_guard: &mut tokio::sync::MutexGuard<'_, KvManager>,
+    signals: &[MoveBlock],
+) -> MoveBlockResponse {
path: lib/llm/src/mocker/scheduler.rs
user: login: jthomson04
body: Like above, use `Result<()>` here
created_at: 2025-05-19T20:11:40Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,586 @@
+// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
path: lib/llm/src/mocker/sequence.rs
user: login: jthomson04
body: Is this all needed? There's some very similar stuff in `crate::tokens`
created_at: 2025-05-19T20:12:20Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,396 @@
+// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
path: lib/llm/src/mocker/kv_manager.rs
user: login: PeaBrane
body: yea, I think that's the eventual goal once the kvbm is stabilized, and / or other frameworks start using it
Also tempted to say to keep this impl, at least for now, for sake of shadowing / mocking the actual kvbm. What's your thought on this?
created_at: 2025-05-19T20:40:09Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,396 @@
+// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use crate::mocker::evictor::LRUEvictor;
+use crate::mocker::protocols::{MoveBlock, MoveBlockResponse, PrefillCost, UniqueBlock};
+use crate::mocker::sequence::ActiveSequence;
+use std::collections::{HashMap, HashSet};
+
+/// Mock implementation of worker for testing and simulation
+pub struct KvManager {
+    max_capacity: usize,
+    block_size: usize,
+    active_blocks: HashMap<UniqueBlock, usize>,
+    inactive_blocks: LRUEvictor<UniqueBlock>,
+    all_blocks: HashSet<UniqueBlock>,
+}
+
+impl KvManager {
+    pub fn new(max_capacity: usize, block_size: usize) -> Self {
+        let active_blocks = HashMap::new();
+        let inactive_blocks = LRUEvictor::new();
+        let all_blocks = HashSet::new();
+
+        KvManager {
+            max_capacity,
+            block_size,
+            active_blocks,
+            inactive_blocks,
+            all_blocks,
+        }
+    }
+
+    /// Get the maximum capacity
+    pub fn max_capacity(&self) -> usize {
+        self.max_capacity
+    }
+
+    /// Get the block size
+    pub fn block_size(&self) -> usize {
+        self.block_size
+    }
+
+    /// Get a reference to the active blocks
+    pub fn active_blocks(&self) -> &HashMap<UniqueBlock, usize> {
+        &self.active_blocks
+    }
+
+    /// Get a reference to the inactive blocks
+    pub fn inactive_blocks(&self) -> &LRUEvictor<UniqueBlock> {
+        &self.inactive_blocks
+    }
+
+    /// Get a reference to all blocks
+    pub fn all_blocks(&self) -> &HashSet<UniqueBlock> {
+        &self.all_blocks
+    }
+
+    /// Process a MoveBlock instruction synchronously
+    pub fn process(&mut self, event: &MoveBlock) -> MoveBlockResponse {
+        match event {
+            MoveBlock::Use(hashes, _) => {
+                for hash in hashes {
+                    // First check if it already exists in active blocks
+                    if let Some(ref_count) = self.active_blocks.get_mut(hash) {
+                        // Block already active, just increment reference count
+                        *ref_count += 1;
+                        continue;
+                    }
+
+                    // Then check if it exists in inactive and move it to active if found
+                    if self.inactive_blocks.remove(hash) {
+                        // Insert into active with reference count 1
+                        self.active_blocks.insert(hash.clone(), 1);
+                        continue;
+                    }
+
+                    // Get counts for capacity check
+                    let active_count = self.active_blocks.len();
+                    let inactive_count = self.inactive_blocks.len();
+
+                    // If at max capacity, evict the oldest entry from inactive blocks
+                    if active_count + inactive_count >= self.max_capacity {
+                        if let Some(evicted) = self.inactive_blocks.evict() {
+                            // Remove evicted block from all_blocks
+                            self.all_blocks.remove(&evicted);
+                        } else {
+                            // Return failure instead of panicking
+                            return MoveBlockResponse::Failure;
+                        }
+                    }
+
+                    // Now insert the new block in active blocks with reference count 1
+                    self.active_blocks.insert(hash.clone(), 1);
+                    // Add to all_blocks as it's a new block
+                    self.all_blocks.insert(hash.clone());
+                }
+            }
+            MoveBlock::Destroy(hashes) => {
+                // Loop in inverse direction
+                for hash in hashes.iter().rev() {
+                    self.active_blocks.remove(hash).unwrap();
+                    // Remove from all_blocks when destroyed
+                    assert!(self.all_blocks.remove(hash));
+                }
+            }
+            MoveBlock::Deref(hashes) => {
+                // Loop in inverse direction
+                for hash in hashes.iter().rev() {
+                    // Decrement reference count and check if we need to move to inactive
+                    if let Some(ref_count) = self.active_blocks.get_mut(hash) {
+                        *ref_count -= 1;
+
+                        // If reference count reaches zero, remove from active and move to inactive
+                        if *ref_count == 0 {
+                            self.active_blocks.remove(hash);
+                            // Use the LRUEvictor's timing functionality
+                            self.inactive_blocks.insert(hash.clone());
+                        }
+                    }
+                }
+            }
+            MoveBlock::Promote(uuid, hash) => {
+                let uuid_block = UniqueBlock::PartialBlock(*uuid);
path: lib/llm/src/mocker/kv_manager.rs
user: login: PeaBrane
body: Yes, I think so. I kinda realized that `tokens.rs` has most that is needed after the fact. May try to interface a bit with it. (Unrelated note, the `kv_router` can probably also use quite a bit from it
created_at: 2025-05-19T20:44:04Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,610 @@
+// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use crate::kv_router::protocols::ForwardPassMetrics;
+use crate::mocker::evictor::LRUEvictor;
+use crate::mocker::kv_manager::KvManager;
+use crate::mocker::protocols::DirectRequest;
+use crate::mocker::protocols::{MoveBlock, MoveBlockResponse, PrefillCost, UniqueBlock};
+use crate::mocker::sequence::ActiveSequence;
+use std::collections::HashMap;
+use std::collections::VecDeque;
+use std::sync::Arc;
+use tokio::sync::{mpsc, Mutex};
+use tokio::task::JoinHandle;
+use tokio::time::{interval, Duration};
+use tokio_util::sync::CancellationToken;
+use uuid::Uuid;
+
+/// Enum representing either a direct request or an active sequence
+pub enum Request {
+    Direct(DirectRequest),
+    Active(ActiveSequence),
+}
+struct SchedulerState {
+    waiting: VecDeque<Uuid>,
+    ready: VecDeque<Uuid>,
+    running: LRUEvictor<Uuid>,
+    requests: HashMap<Uuid, Request>,
+    prefill_costs: HashMap<Uuid, Option<PrefillCost>>,
+}
+
+impl Default for SchedulerState {
+    fn default() -> Self {
+        Self {
+            waiting: VecDeque::new(),
+            ready: VecDeque::new(),
+            running: LRUEvictor::new(),
+            requests: HashMap::new(),
+            prefill_costs: HashMap::new(),
+        }
+    }
+}
+
+impl SchedulerState {
+    /// Create a new UUID for a DirectRequest, add it to requests, and push the UUID to waiting.
+    fn receive(&mut self, request: DirectRequest) -> Uuid {
+        // Use the provided UUID if available, otherwise generate a new one
+        let uuid = request.uuid.unwrap_or_else(Uuid::new_v4);
+
+        // Add the request to the map and waiting queue
+        self.requests.insert(uuid, Request::Direct(request));
+        self.waiting.push_back(uuid);
+        uuid
+    }
+
+    /// Get the next UUID from ready or waiting queue and its associated Request.
+    /// Returns from ready if not empty, otherwise from waiting, or None if both are empty.
+    /// Also removes the Request from the requests HashMap.
+    fn next(&mut self) -> Option<(Uuid, Request)> {
+        let uuid = self
+            .ready
+            .pop_front()
+            .or_else(|| self.waiting.pop_front())?;
+        let request = self.requests.remove(&uuid)?;
+        Some((uuid, request))
+    }
+
+    /// Move a UUID and its Request to the ready queue.
+    fn make_ready(&mut self, uuid: Uuid, active_seq: ActiveSequence) {
+        self.requests.insert(uuid, Request::Active(active_seq));
+        self.ready.push_back(uuid);
+    }
+
+    /// Schedule the request with the given UUID.
+    /// Returns the creation signal from the ActiveSequence.
+    fn run(&mut self, uuid: Uuid, active_seq: ActiveSequence) -> MoveBlock {
+        // Insert the request into the map
+        self.requests.insert(uuid, Request::Active(active_seq));
+
+        // Get the creation signal
+        let Some(Request::Active(sequence)) = self.requests.get(&uuid) else {
+            panic!("Failed to get ActiveSequence for UUID");
+        };
+        let Some(signal) = sequence.creation_signal() else {
+            panic!("Failed to get creation signal from ActiveSequence");
+        };
+
+        // Add to running requests
+        self.running.insert(uuid);
+        signal.clone()
+    }
+
+    /// Set the prefill cost for a UUID
+    fn set_prefill_cost(&mut self, uuid: Uuid, cost: Option<PrefillCost>) {
+        self.prefill_costs.insert(uuid, cost);
+    }
+
+    /// Get the prefill compute value for a UUID if available
+    fn get_prefill_compute(&self, uuid: &Uuid) -> Option<f64> {
+        self.prefill_costs
+            .get(uuid)
+            .and_then(|cost| cost.as_ref())
+            .map(|cost| cost.prefill_compute)
+    }
+
+    /// Calculate the current running batched tokens
+    fn num_batched_tokens(&self) -> usize {
+        self.prefill_costs
+            .values()
+            .map(|cost| match cost {
+                Some(cost) => cost.new_tokens,
+                None => 1,
+            })
+            .sum()
+    }
+
+    /// Remove a UUID and its associated Request from collections.
+    fn complete(&mut self, uuid: &Uuid) {
+        self.running.remove(uuid);
+        self.requests.remove(uuid);
+        self.prefill_costs.remove(uuid);
+    }
+
+    /// Preempt the oldest running request by evicting it from running, resetting the sequence,
+    /// and adding it back to the waiting queue.
+    /// Returns the signal from reset_with_signal or None if no requests are running.
+    fn preempt(&mut self) -> Option<Vec<MoveBlock>> {
+        // Evict the oldest UUID from running
+        let uuid = self.running.evict()?;
+        eprintln!("Request {} will be preempted", uuid);
+
+        // Remove the request from the requests HashMap and ensure it's an ActiveSequence
+        let request = self.requests.remove(&uuid)?;
+
+        // Remove the prefill cost to force recomputation
+        self.prefill_costs.remove(&uuid);
+
+        // Extract the ActiveSequence from the Request enum
+        let Request::Active(active_sequence) = request else {
+            panic!("Expected ActiveSequence in running queue")
+        };
+
+        // Reset the sequence and get the new sequence and signal
+        let (new_sequence, signals) = active_sequence.reset_with_signal();
+
+        // Insert the new sequence back into the requests map and add to waiting queue
+        self.requests.insert(uuid, Request::Active(new_sequence));
+        self.waiting.push_back(uuid);
+
+        Some(signals)
+    }
+}
+
+/// Manages scheduling of requests using KvManager resources
+pub struct Scheduler {
+    state: Arc<Mutex<SchedulerState>>,
+    kv_manager: Arc<Mutex<KvManager>>,
+    token_capacity: usize,
+    watermark: f64,
+    block_size: usize,
+    chunk_size: usize,
+    background_handle: Option<JoinHandle<()>>,
+    request_tx: mpsc::Sender<DirectRequest>,
+    cancellation_token: CancellationToken,
+}
+
+impl Scheduler {
+    /// Create a new Scheduler with the given parameters
+    pub fn new(
+        kv_capacity: usize,
+        watermark: f64,
+        block_size: usize,
+        chunk_size: Option<usize>,
+        output_tx: Option<mpsc::Sender<Uuid>>,
+        cancellation_token: Option<CancellationToken>,
+    ) -> Self {
+        // Create KvManager internally
+        let kv_manager = KvManager::new(kv_capacity, block_size);
+
+        let token_capacity: usize = 8192;
+        let state = Arc::new(Mutex::new(SchedulerState::default()));
path: lib/llm/src/mocker/scheduler.rs
user: login: PeaBrane
body: In my initial impl, the `SchedulerState` was quite a bit simpler, but it made the `Scheduler` quite spaghetti, so thought it'd make sense to let the state handle some of the logic. 
In this case here, there does not seem to be any write lock contention (all write arms are never concurrent), so the only overhead may be on the `read / metrics` side. Tempted to say `Mutex` is the to-go but I may be wrong.
Another thing I worry about is deadlocks if we don't "aggregate" these states, in the case we do decide to do concurrent writes in the future, and need to worry about lock acquisition orders. Of course, we're giving up some performance
created_at: 2025-05-19T20:50:18Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,586 @@
+// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
path: lib/llm/src/mocker/sequence.rs
user: login: PeaBrane
body: moving conversation to above thread
created_at: 2025-05-19T20:53:19Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,396 @@
+// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use crate::mocker::evictor::LRUEvictor;
+use crate::mocker::protocols::{MoveBlock, MoveBlockResponse, PrefillCost, UniqueBlock};
+use crate::mocker::sequence::ActiveSequence;
+use std::collections::{HashMap, HashSet};
+
+/// Mock implementation of worker for testing and simulation
+pub struct KvManager {
+    max_capacity: usize,
+    block_size: usize,
+    active_blocks: HashMap<UniqueBlock, usize>,
+    inactive_blocks: LRUEvictor<UniqueBlock>,
+    all_blocks: HashSet<UniqueBlock>,
+}
+
+impl KvManager {
+    pub fn new(max_capacity: usize, block_size: usize) -> Self {
+        let active_blocks = HashMap::new();
+        let inactive_blocks = LRUEvictor::new();
+        let all_blocks = HashSet::new();
+
+        KvManager {
+            max_capacity,
+            block_size,
+            active_blocks,
+            inactive_blocks,
+            all_blocks,
+        }
+    }
+
+    /// Get the maximum capacity
+    pub fn max_capacity(&self) -> usize {
+        self.max_capacity
+    }
+
+    /// Get the block size
+    pub fn block_size(&self) -> usize {
+        self.block_size
+    }
+
+    /// Get a reference to the active blocks
+    pub fn active_blocks(&self) -> &HashMap<UniqueBlock, usize> {
+        &self.active_blocks
+    }
+
+    /// Get a reference to the inactive blocks
+    pub fn inactive_blocks(&self) -> &LRUEvictor<UniqueBlock> {
+        &self.inactive_blocks
+    }
+
+    /// Get a reference to all blocks
+    pub fn all_blocks(&self) -> &HashSet<UniqueBlock> {
+        &self.all_blocks
+    }
+
+    /// Process a MoveBlock instruction synchronously
+    pub fn process(&mut self, event: &MoveBlock) -> MoveBlockResponse {
+        match event {
+            MoveBlock::Use(hashes, _) => {
+                for hash in hashes {
+                    // First check if it already exists in active blocks
+                    if let Some(ref_count) = self.active_blocks.get_mut(hash) {
+                        // Block already active, just increment reference count
+                        *ref_count += 1;
+                        continue;
+                    }
+
+                    // Then check if it exists in inactive and move it to active if found
+                    if self.inactive_blocks.remove(hash) {
+                        // Insert into active with reference count 1
+                        self.active_blocks.insert(hash.clone(), 1);
+                        continue;
+                    }
+
+                    // Get counts for capacity check
+                    let active_count = self.active_blocks.len();
+                    let inactive_count = self.inactive_blocks.len();
+
+                    // If at max capacity, evict the oldest entry from inactive blocks
+                    if active_count + inactive_count >= self.max_capacity {
+                        if let Some(evicted) = self.inactive_blocks.evict() {
+                            // Remove evicted block from all_blocks
+                            self.all_blocks.remove(&evicted);
+                        } else {
+                            // Return failure instead of panicking
+                            return MoveBlockResponse::Failure;
+                        }
+                    }
+
+                    // Now insert the new block in active blocks with reference count 1
+                    self.active_blocks.insert(hash.clone(), 1);
+                    // Add to all_blocks as it's a new block
+                    self.all_blocks.insert(hash.clone());
+                }
+            }
+            MoveBlock::Destroy(hashes) => {
+                // Loop in inverse direction
+                for hash in hashes.iter().rev() {
+                    self.active_blocks.remove(hash).unwrap();
+                    // Remove from all_blocks when destroyed
+                    assert!(self.all_blocks.remove(hash));
+                }
+            }
+            MoveBlock::Deref(hashes) => {
+                // Loop in inverse direction
+                for hash in hashes.iter().rev() {
path: lib/llm/src/mocker/kv_manager.rs
user: login: PeaBrane
body: I'll try to work in something like that, but probably `Rc` for now, single-threaded. One thing I worry about is if we expose the `active_blocks` via a clone in a "getter" method, it may screw the reference count up, but then again we should probably not expose that in the first place 
created_at: 2025-05-19T21:08:14Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,610 @@
+// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use crate::kv_router::protocols::ForwardPassMetrics;
+use crate::mocker::evictor::LRUEvictor;
+use crate::mocker::kv_manager::KvManager;
+use crate::mocker::protocols::DirectRequest;
+use crate::mocker::protocols::{MoveBlock, MoveBlockResponse, PrefillCost, UniqueBlock};
+use crate::mocker::sequence::ActiveSequence;
+use std::collections::HashMap;
+use std::collections::VecDeque;
+use std::sync::Arc;
+use tokio::sync::{mpsc, Mutex};
+use tokio::task::JoinHandle;
+use tokio::time::{interval, Duration};
+use tokio_util::sync::CancellationToken;
+use uuid::Uuid;
+
+/// Enum representing either a direct request or an active sequence
+pub enum Request {
+    Direct(DirectRequest),
+    Active(ActiveSequence),
+}
+struct SchedulerState {
+    waiting: VecDeque<Uuid>,
+    ready: VecDeque<Uuid>,
+    running: LRUEvictor<Uuid>,
+    requests: HashMap<Uuid, Request>,
+    prefill_costs: HashMap<Uuid, Option<PrefillCost>>,
+}
+
+impl Default for SchedulerState {
+    fn default() -> Self {
+        Self {
+            waiting: VecDeque::new(),
+            ready: VecDeque::new(),
+            running: LRUEvictor::new(),
+            requests: HashMap::new(),
+            prefill_costs: HashMap::new(),
+        }
+    }
+}
+
+impl SchedulerState {
+    /// Create a new UUID for a DirectRequest, add it to requests, and push the UUID to waiting.
+    fn receive(&mut self, request: DirectRequest) -> Uuid {
+        // Use the provided UUID if available, otherwise generate a new one
+        let uuid = request.uuid.unwrap_or_else(Uuid::new_v4);
+
+        // Add the request to the map and waiting queue
+        self.requests.insert(uuid, Request::Direct(request));
+        self.waiting.push_back(uuid);
+        uuid
+    }
+
+    /// Get the next UUID from ready or waiting queue and its associated Request.
+    /// Returns from ready if not empty, otherwise from waiting, or None if both are empty.
+    /// Also removes the Request from the requests HashMap.
+    fn next(&mut self) -> Option<(Uuid, Request)> {
+        let uuid = self
+            .ready
+            .pop_front()
+            .or_else(|| self.waiting.pop_front())?;
+        let request = self.requests.remove(&uuid)?;
+        Some((uuid, request))
+    }
+
+    /// Move a UUID and its Request to the ready queue.
+    fn make_ready(&mut self, uuid: Uuid, active_seq: ActiveSequence) {
+        self.requests.insert(uuid, Request::Active(active_seq));
+        self.ready.push_back(uuid);
+    }
+
+    /// Schedule the request with the given UUID.
+    /// Returns the creation signal from the ActiveSequence.
+    fn run(&mut self, uuid: Uuid, active_seq: ActiveSequence) -> MoveBlock {
+        // Insert the request into the map
+        self.requests.insert(uuid, Request::Active(active_seq));
+
+        // Get the creation signal
+        let Some(Request::Active(sequence)) = self.requests.get(&uuid) else {
+            panic!("Failed to get ActiveSequence for UUID");
+        };
+        let Some(signal) = sequence.creation_signal() else {
+            panic!("Failed to get creation signal from ActiveSequence");
+        };
+
+        // Add to running requests
+        self.running.insert(uuid);
+        signal.clone()
+    }
+
+    /// Set the prefill cost for a UUID
+    fn set_prefill_cost(&mut self, uuid: Uuid, cost: Option<PrefillCost>) {
+        self.prefill_costs.insert(uuid, cost);
+    }
+
+    /// Get the prefill compute value for a UUID if available
+    fn get_prefill_compute(&self, uuid: &Uuid) -> Option<f64> {
+        self.prefill_costs
+            .get(uuid)
+            .and_then(|cost| cost.as_ref())
+            .map(|cost| cost.prefill_compute)
+    }
+
+    /// Calculate the current running batched tokens
+    fn num_batched_tokens(&self) -> usize {
+        self.prefill_costs
+            .values()
+            .map(|cost| match cost {
+                Some(cost) => cost.new_tokens,
+                None => 1,
+            })
+            .sum()
+    }
+
+    /// Remove a UUID and its associated Request from collections.
+    fn complete(&mut self, uuid: &Uuid) {
+        self.running.remove(uuid);
+        self.requests.remove(uuid);
+        self.prefill_costs.remove(uuid);
+    }
+
+    /// Preempt the oldest running request by evicting it from running, resetting the sequence,
+    /// and adding it back to the waiting queue.
+    /// Returns the signal from reset_with_signal or None if no requests are running.
+    fn preempt(&mut self) -> Option<Vec<MoveBlock>> {
+        // Evict the oldest UUID from running
+        let uuid = self.running.evict()?;
+        eprintln!("Request {} will be preempted", uuid);
+
+        // Remove the request from the requests HashMap and ensure it's an ActiveSequence
+        let request = self.requests.remove(&uuid)?;
+
+        // Remove the prefill cost to force recomputation
+        self.prefill_costs.remove(&uuid);
+
+        // Extract the ActiveSequence from the Request enum
+        let Request::Active(active_sequence) = request else {
+            panic!("Expected ActiveSequence in running queue")
+        };
+
+        // Reset the sequence and get the new sequence and signal
+        let (new_sequence, signals) = active_sequence.reset_with_signal();
+
+        // Insert the new sequence back into the requests map and add to waiting queue
+        self.requests.insert(uuid, Request::Active(new_sequence));
+        self.waiting.push_back(uuid);
+
+        Some(signals)
+    }
+}
+
+/// Manages scheduling of requests using KvManager resources
+pub struct Scheduler {
+    state: Arc<Mutex<SchedulerState>>,
+    kv_manager: Arc<Mutex<KvManager>>,
+    token_capacity: usize,
+    watermark: f64,
+    block_size: usize,
+    chunk_size: usize,
+    background_handle: Option<JoinHandle<()>>,
+    request_tx: mpsc::Sender<DirectRequest>,
+    cancellation_token: CancellationToken,
+}
+
+impl Scheduler {
+    /// Create a new Scheduler with the given parameters
+    pub fn new(
+        kv_capacity: usize,
+        watermark: f64,
+        block_size: usize,
+        chunk_size: Option<usize>,
+        output_tx: Option<mpsc::Sender<Uuid>>,
+        cancellation_token: Option<CancellationToken>,
+    ) -> Self {
+        // Create KvManager internally
+        let kv_manager = KvManager::new(kv_capacity, block_size);
+
+        let token_capacity: usize = 8192;
+        let state = Arc::new(Mutex::new(SchedulerState::default()));
+
+        let kv_manager = Arc::new(Mutex::new(kv_manager));
+        let chunk_size = chunk_size.unwrap_or(256);
+
+        // Create channel for request handling
+        let (request_tx, mut request_rx) = mpsc::channel::<DirectRequest>(1024);
+
+        // Use provided cancellation token or create new one
+        let cancellation_token = cancellation_token.unwrap_or_default();
+        let token_clone = cancellation_token.clone();
+
+        // Create a clone for the background task
+        let state_clone = state.clone();
+        let kv_manager_clone = kv_manager.clone();
+        let output_tx_clone = output_tx.clone();
+
+        // Spawn main background task with cancellation token
+        let background_handle = tokio::spawn(async move {
+            let mut schedule_interval = interval(Duration::from_millis(5));
+            let mut simulate_interval = interval(Duration::from_millis(1));
+
+            loop {
+                tokio::select! {
+                    biased;
+
+                    // Enqueue new request
+                    Some(request) = request_rx.recv() => {
+                        let mut state = state_clone.lock().await;
+                        state.receive(request);
+                    }
+
+                    // Try Scheduling Requests
+                    _ = schedule_interval.tick() => {
+                        let mut state_guard = state_clone.lock().await;
+                        let mut kv_manager_guard = kv_manager_clone.lock().await;
+
+                        // Process DirectRequests, converting them to ActiveSequence and scheduling them until we can't
+                        // schedule anymore.
+                        while let Some((uuid, request)) = state_guard.next() {
+                            let active_sequence = get_active_sequence(request, block_size, chunk_size);
+
+                            // Calculate token budget using new_tokens from PrefillCost
+                            let total_prefill_tokens = state_guard.num_batched_tokens();
+                            let tokens_budget = token_capacity.saturating_sub(total_prefill_tokens);
+
+                            // Check if it can be scheduled
+                            let Some(prefill_cost) = kv_manager_guard.try_schedule(&active_sequence, watermark, tokens_budget) else {
+                                state_guard.make_ready(uuid, active_sequence);
+                                break;
+                            };
+
+                            // Get creation signal and schedule the request
+                            let signal = state_guard.run(uuid, active_sequence);
+                            kv_manager_guard.process(&signal);
+                            state_guard.set_prefill_cost(uuid, Some(prefill_cost));
+                        }
+                    }
+
+                    // Check for cancellation
+                    _ = token_clone.cancelled() => {
+                        break;
+                    }
+
+                    // Simulate running requests (prefill + decode)
+                    _ = simulate_interval.tick() => {
+                        let mut state_guard = state_clone.lock().await;
+                        let mut kv_manager_guard = kv_manager_clone.lock().await;
+
+                        // Process each running request
+                        let mut total_sleep_duration = Duration::from_millis(0);
+                        let uuids: Vec<Uuid> = state_guard.running.keys().cloned().collect();
+                        for uuid in uuids {
+                            // Check if UUID is still in running_requests, if not skip this iteration
+                            if !state_guard.running.contains(&uuid) {
+                                continue;
+                            }
+
+                            // Get prefill compute value first
+                            let prefill_compute = state_guard.get_prefill_compute(&uuid);
+
+                            // Get the active sequence for this UUID
+                            let sequence = state_guard.requests.get_mut(&uuid)
+                                .and_then(|req| if let Request::Active(seq) = req { Some(seq) } else { None })
+                                .expect("UUID in running_requests must have a corresponding active sequence");
+
+                            // Generate token and get signals
+                            let signals = sequence.generate();
+
+                            // Accumulate sleep duration based on prefill_compute if available
+                            let sleep_ms = if let Some(compute) = prefill_compute {
+                                (compute / 131072.0) as u64  // TODO: A magic scaling to model prefill compute
+                            } else {
+                                1  // FIXME: A dummy number to simulate ITL. Need some work.
+                            };
+                            total_sleep_duration += Duration::from_millis(sleep_ms);
+
+                            // Process all signals with the KvManager
+                            if process_signals(&mut kv_manager_guard, &signals) == MoveBlockResponse::Failure {
+                                sequence.partial_reset();
+
+                                // free_signal derefs the preempted blocks
+                                let Some(free_signal) = state_guard.preempt() else {
+                                    panic!("Failed to acquire signal to free KV blocks from preemption");
+                                };
+
+                                for signal in free_signal {
+                                    kv_manager_guard.process(&signal);
+                                }
+                                continue;
+                            }
+
+                            // Send UUID notification for each generated token
+                            // TODO: hook this up to an AsyncEngine
+                            if let Some(tx) = &output_tx_clone {
+                                let _ = tx.try_send(uuid);
+                            }
+
+                            // Check if we're done after generating
+                            if sequence.generated_tokens() >= sequence.max_output_tokens() {
+                                state_guard.complete(&uuid);
+                                continue;
+                            }
+
+                            // Transition to decode (no prefill cost)
+                            if sequence.generated_tokens() == 1 {
+                                state_guard.set_prefill_cost(uuid, None);
+                            }
+                        }
+
+                        // Sleep once for the accumulated duration
+                        if total_sleep_duration.as_millis() > 0 {
+                            tokio::time::sleep(total_sleep_duration).await;
+                        }
+                    }
+                }
+            }
+        });
+
+        Self {
+            state,
+            kv_manager,
+            token_capacity,
+            watermark,
+            block_size,
+            chunk_size,
+            background_handle: Some(background_handle),
+            request_tx,
+            cancellation_token,
+        }
+    }
+
+    /// Add a new request to the waiting queue
+    pub async fn receive(&self, request: DirectRequest) {
+        let _ = self.request_tx.send(request).await;
+    }
+
+    /// Get the count of waiting requests
+    pub async fn waiting_count(&self) -> usize {
+        let state = self.state.lock().await;
+        state.waiting.len()
+    }
+
+    /// Get the count of running requests
+    pub async fn running_count(&self) -> usize {
+        let state = self.state.lock().await;
+        state.running.len()
+    }
+
+    /// Get the current capacity of the KvManager
+    pub async fn kv_usage_perc(&self) -> f64 {
+        let kv_manager = self.kv_manager.lock().await;
+        kv_manager.current_capacity_perc()
+    }
+
+    /// Returns forward pass metrics for monitoring purposes
+    pub async fn get_forward_pass_metrics(&self) -> ForwardPassMetrics {
+        let state = self.state.lock().await;
+        let kv_manager = self.kv_manager.lock().await;
+
+        // Get the active blocks and total capacity from KvManager
+        let active_blocks_count = kv_manager.active_blocks().len() as u64;
+        let total_capacity = kv_manager.max_capacity() as u64;
+
+        // Calculate GPU cache usage percentage
+        let gpu_cache_usage_perc = if total_capacity > 0 {
+            active_blocks_count as f32 / total_capacity as f32
+        } else {
+            0.0
+        };
+
+        ForwardPassMetrics {
+            request_active_slots: state.running.len() as u64,
+            request_total_slots: 420, // Dummy value as specified
+            kv_active_blocks: active_blocks_count,
+            kv_total_blocks: total_capacity,
+            num_requests_waiting: state.waiting.len() as u64,
+            gpu_cache_usage_perc,
+            gpu_prefix_cache_hit_rate: 0.0, // Placeholder value as specified
+        }
+    }
+}
+
+// Implement Clone for Scheduler to support sharing between tasks
+impl Clone for Scheduler {
+    fn clone(&self) -> Self {
+        Self {
+            state: self.state.clone(),
+            kv_manager: self.kv_manager.clone(),
+            token_capacity: self.token_capacity,
+            watermark: self.watermark,
+            block_size: self.block_size,
+            chunk_size: self.chunk_size,
+            background_handle: None,
+            request_tx: self.request_tx.clone(),
+            cancellation_token: self.cancellation_token.clone(),
+        }
+    }
+}
+
+impl Drop for Scheduler {
+    fn drop(&mut self) {
+        // Only the original instance should cancel
+        if self.background_handle.is_some() {
+            self.cancellation_token.cancel();
+        }
+
+        // Abort the background task if this is the original instance
+        if let Some(handle) = self.background_handle.take() {
+            handle.abort();
+        }
+    }
+}
+
+/// Convert a Request to an ActiveSequence
+fn get_active_sequence(request: Request, block_size: usize, chunk_size: usize) -> ActiveSequence {
+    if let Request::Active(active_seq) = request {
+        return active_seq;
+    }
+
+    let Request::Direct(direct_request) = request else {
+        unreachable!("Request must be either Direct or Active");
+    };
+
+    ActiveSequence::new(
+        direct_request.tokens,
+        direct_request.max_output_tokens,
+        Some(block_size),
+        Some(chunk_size),
+    )
+}
+
+/// Processes MoveBlock signals with the KvManager.
+///
+/// When a signal fails, this function verifies that the failure is for an expected case:
+/// specifically a single signal attempting to create a single partial (generation) block.
+/// This validation is important because in normal operation, the only legitimate failure
+/// case should be when trying to acquire a new generation block - any other failures would
+/// indicate an unexpected state in the system.
+fn process_signals(
+    kv_manager_guard: &mut tokio::sync::MutexGuard<'_, KvManager>,
+    signals: &[MoveBlock],
+) -> MoveBlockResponse {
path: lib/llm/src/mocker/scheduler.rs
user: login: PeaBrane
body: opted for a bool in the end
created_at: 2025-05-19T21:47:46Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,193 @@
+// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::cmp::Eq;
+use std::collections::{HashMap, VecDeque};
+use std::hash::Hash;
+use std::time::Instant;
+
+/// An LRU evictor that maintains objects and evicts them based on their
+/// last accessed time. Implements a "lazy" eviction mechanism where:
+/// 1. The priority queue does not immediately reflect updates or removes
+/// 2. Objects are pushed to the queue in order of increasing priority (older objects first)
+/// 3. The user must ensure objects are added in correct priority (temporal order)
+/// 4. Remove and update operations are lazy - entries remain in the queue until
+///    they are either evicted or cleaned up during maintenance
+#[derive(Debug)]
+pub struct LRUEvictor<T: Clone + Eq + Hash> {
+    free_table: HashMap<T, f64>,
+    priority_queue: VecDeque<(T, f64)>,
+    cleanup_threshold: usize,
+    start_time: Instant,
+}
+
+impl<T: Clone + Eq + Hash> Default for LRUEvictor<T> {
path: lib/llm/src/mocker/evictor.rs
user: login: PeaBrane
body: for some reason clippy wasn't happy with not having a default here, but reorganized the defaults a bit to make things cleaner across the board
created_at: 2025-05-19T22:11:46Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,610 @@
+// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use crate::kv_router::protocols::ForwardPassMetrics;
+use crate::mocker::evictor::LRUEvictor;
+use crate::mocker::kv_manager::KvManager;
+use crate::mocker::protocols::DirectRequest;
+use crate::mocker::protocols::{MoveBlock, MoveBlockResponse, PrefillCost, UniqueBlock};
+use crate::mocker::sequence::ActiveSequence;
+use std::collections::HashMap;
+use std::collections::VecDeque;
+use std::sync::Arc;
+use tokio::sync::{mpsc, Mutex};
+use tokio::task::JoinHandle;
+use tokio::time::{interval, Duration};
+use tokio_util::sync::CancellationToken;
+use uuid::Uuid;
+
+/// Enum representing either a direct request or an active sequence
+pub enum Request {
+    Direct(DirectRequest),
+    Active(ActiveSequence),
+}
+struct SchedulerState {
+    waiting: VecDeque<Uuid>,
+    ready: VecDeque<Uuid>,
+    running: LRUEvictor<Uuid>,
+    requests: HashMap<Uuid, Request>,
+    prefill_costs: HashMap<Uuid, Option<PrefillCost>>,
+}
+
+impl Default for SchedulerState {
+    fn default() -> Self {
+        Self {
+            waiting: VecDeque::new(),
+            ready: VecDeque::new(),
+            running: LRUEvictor::new(),
+            requests: HashMap::new(),
+            prefill_costs: HashMap::new(),
+        }
+    }
+}
+
+impl SchedulerState {
+    /// Create a new UUID for a DirectRequest, add it to requests, and push the UUID to waiting.
+    fn receive(&mut self, request: DirectRequest) -> Uuid {
+        // Use the provided UUID if available, otherwise generate a new one
+        let uuid = request.uuid.unwrap_or_else(Uuid::new_v4);
+
+        // Add the request to the map and waiting queue
+        self.requests.insert(uuid, Request::Direct(request));
+        self.waiting.push_back(uuid);
+        uuid
+    }
+
+    /// Get the next UUID from ready or waiting queue and its associated Request.
+    /// Returns from ready if not empty, otherwise from waiting, or None if both are empty.
+    /// Also removes the Request from the requests HashMap.
+    fn next(&mut self) -> Option<(Uuid, Request)> {
+        let uuid = self
+            .ready
+            .pop_front()
+            .or_else(|| self.waiting.pop_front())?;
+        let request = self.requests.remove(&uuid)?;
+        Some((uuid, request))
+    }
+
+    /// Move a UUID and its Request to the ready queue.
+    fn make_ready(&mut self, uuid: Uuid, active_seq: ActiveSequence) {
+        self.requests.insert(uuid, Request::Active(active_seq));
+        self.ready.push_back(uuid);
+    }
+
+    /// Schedule the request with the given UUID.
+    /// Returns the creation signal from the ActiveSequence.
+    fn run(&mut self, uuid: Uuid, active_seq: ActiveSequence) -> MoveBlock {
+        // Insert the request into the map
+        self.requests.insert(uuid, Request::Active(active_seq));
+
+        // Get the creation signal
+        let Some(Request::Active(sequence)) = self.requests.get(&uuid) else {
+            panic!("Failed to get ActiveSequence for UUID");
+        };
+        let Some(signal) = sequence.creation_signal() else {
+            panic!("Failed to get creation signal from ActiveSequence");
+        };
+
+        // Add to running requests
+        self.running.insert(uuid);
+        signal.clone()
+    }
+
+    /// Set the prefill cost for a UUID
+    fn set_prefill_cost(&mut self, uuid: Uuid, cost: Option<PrefillCost>) {
+        self.prefill_costs.insert(uuid, cost);
+    }
+
+    /// Get the prefill compute value for a UUID if available
+    fn get_prefill_compute(&self, uuid: &Uuid) -> Option<f64> {
+        self.prefill_costs
+            .get(uuid)
+            .and_then(|cost| cost.as_ref())
+            .map(|cost| cost.prefill_compute)
+    }
+
+    /// Calculate the current running batched tokens
+    fn num_batched_tokens(&self) -> usize {
+        self.prefill_costs
+            .values()
+            .map(|cost| match cost {
+                Some(cost) => cost.new_tokens,
+                None => 1,
+            })
+            .sum()
+    }
+
+    /// Remove a UUID and its associated Request from collections.
+    fn complete(&mut self, uuid: &Uuid) {
+        self.running.remove(uuid);
+        self.requests.remove(uuid);
+        self.prefill_costs.remove(uuid);
+    }
+
+    /// Preempt the oldest running request by evicting it from running, resetting the sequence,
+    /// and adding it back to the waiting queue.
+    /// Returns the signal from reset_with_signal or None if no requests are running.
+    fn preempt(&mut self) -> Option<Vec<MoveBlock>> {
+        // Evict the oldest UUID from running
+        let uuid = self.running.evict()?;
+        eprintln!("Request {} will be preempted", uuid);
+
+        // Remove the request from the requests HashMap and ensure it's an ActiveSequence
+        let request = self.requests.remove(&uuid)?;
+
+        // Remove the prefill cost to force recomputation
+        self.prefill_costs.remove(&uuid);
+
+        // Extract the ActiveSequence from the Request enum
+        let Request::Active(active_sequence) = request else {
+            panic!("Expected ActiveSequence in running queue")
+        };
+
+        // Reset the sequence and get the new sequence and signal
+        let (new_sequence, signals) = active_sequence.reset_with_signal();
+
+        // Insert the new sequence back into the requests map and add to waiting queue
+        self.requests.insert(uuid, Request::Active(new_sequence));
+        self.waiting.push_back(uuid);
+
+        Some(signals)
+    }
+}
+
+/// Manages scheduling of requests using KvManager resources
+pub struct Scheduler {
+    state: Arc<Mutex<SchedulerState>>,
+    kv_manager: Arc<Mutex<KvManager>>,
+    token_capacity: usize,
+    watermark: f64,
+    block_size: usize,
+    chunk_size: usize,
+    background_handle: Option<JoinHandle<()>>,
+    request_tx: mpsc::Sender<DirectRequest>,
+    cancellation_token: CancellationToken,
+}
+
+impl Scheduler {
+    /// Create a new Scheduler with the given parameters
+    pub fn new(
+        kv_capacity: usize,
+        watermark: f64,
+        block_size: usize,
+        chunk_size: Option<usize>,
+        output_tx: Option<mpsc::Sender<Uuid>>,
+        cancellation_token: Option<CancellationToken>,
+    ) -> Self {
+        // Create KvManager internally
+        let kv_manager = KvManager::new(kv_capacity, block_size);
+
+        let token_capacity: usize = 8192;
+        let state = Arc::new(Mutex::new(SchedulerState::default()));
+
+        let kv_manager = Arc::new(Mutex::new(kv_manager));
+        let chunk_size = chunk_size.unwrap_or(256);
+
+        // Create channel for request handling
+        let (request_tx, mut request_rx) = mpsc::channel::<DirectRequest>(1024);
+
+        // Use provided cancellation token or create new one
+        let cancellation_token = cancellation_token.unwrap_or_default();
+        let token_clone = cancellation_token.clone();
+
+        // Create a clone for the background task
+        let state_clone = state.clone();
+        let kv_manager_clone = kv_manager.clone();
+        let output_tx_clone = output_tx.clone();
+
+        // Spawn main background task with cancellation token
+        let background_handle = tokio::spawn(async move {
+            let mut schedule_interval = interval(Duration::from_millis(5));
+            let mut simulate_interval = interval(Duration::from_millis(1));
+
+            loop {
+                tokio::select! {
+                    biased;
+
+                    // Enqueue new request
+                    Some(request) = request_rx.recv() => {
+                        let mut state = state_clone.lock().await;
+                        state.receive(request);
+                    }
+
+                    // Try Scheduling Requests
+                    _ = schedule_interval.tick() => {
+                        let mut state_guard = state_clone.lock().await;
+                        let mut kv_manager_guard = kv_manager_clone.lock().await;
+
+                        // Process DirectRequests, converting them to ActiveSequence and scheduling them until we can't
+                        // schedule anymore.
+                        while let Some((uuid, request)) = state_guard.next() {
+                            let active_sequence = get_active_sequence(request, block_size, chunk_size);
+
+                            // Calculate token budget using new_tokens from PrefillCost
+                            let total_prefill_tokens = state_guard.num_batched_tokens();
+                            let tokens_budget = token_capacity.saturating_sub(total_prefill_tokens);
+
+                            // Check if it can be scheduled
+                            let Some(prefill_cost) = kv_manager_guard.try_schedule(&active_sequence, watermark, tokens_budget) else {
+                                state_guard.make_ready(uuid, active_sequence);
+                                break;
+                            };
+
+                            // Get creation signal and schedule the request
+                            let signal = state_guard.run(uuid, active_sequence);
+                            kv_manager_guard.process(&signal);
+                            state_guard.set_prefill_cost(uuid, Some(prefill_cost));
+                        }
+                    }
+
+                    // Check for cancellation
+                    _ = token_clone.cancelled() => {
+                        break;
+                    }
+
+                    // Simulate running requests (prefill + decode)
+                    _ = simulate_interval.tick() => {
+                        let mut state_guard = state_clone.lock().await;
+                        let mut kv_manager_guard = kv_manager_clone.lock().await;
+
+                        // Process each running request
+                        let mut total_sleep_duration = Duration::from_millis(0);
+                        let uuids: Vec<Uuid> = state_guard.running.keys().cloned().collect();
+                        for uuid in uuids {
+                            // Check if UUID is still in running_requests, if not skip this iteration
+                            if !state_guard.running.contains(&uuid) {
+                                continue;
+                            }
+
+                            // Get prefill compute value first
+                            let prefill_compute = state_guard.get_prefill_compute(&uuid);
+
+                            // Get the active sequence for this UUID
+                            let sequence = state_guard.requests.get_mut(&uuid)
+                                .and_then(|req| if let Request::Active(seq) = req { Some(seq) } else { None })
+                                .expect("UUID in running_requests must have a corresponding active sequence");
+
+                            // Generate token and get signals
+                            let signals = sequence.generate();
+
+                            // Accumulate sleep duration based on prefill_compute if available
+                            let sleep_ms = if let Some(compute) = prefill_compute {
+                                (compute / 131072.0) as u64  // TODO: A magic scaling to model prefill compute
+                            } else {
+                                1  // FIXME: A dummy number to simulate ITL. Need some work.
+                            };
+                            total_sleep_duration += Duration::from_millis(sleep_ms);
+
+                            // Process all signals with the KvManager
+                            if process_signals(&mut kv_manager_guard, &signals) == MoveBlockResponse::Failure {
+                                sequence.partial_reset();
+
+                                // free_signal derefs the preempted blocks
+                                let Some(free_signal) = state_guard.preempt() else {
+                                    panic!("Failed to acquire signal to free KV blocks from preemption");
+                                };
+
+                                for signal in free_signal {
+                                    kv_manager_guard.process(&signal);
+                                }
+                                continue;
+                            }
+
+                            // Send UUID notification for each generated token
+                            // TODO: hook this up to an AsyncEngine
+                            if let Some(tx) = &output_tx_clone {
+                                let _ = tx.try_send(uuid);
+                            }
+
+                            // Check if we're done after generating
+                            if sequence.generated_tokens() >= sequence.max_output_tokens() {
+                                state_guard.complete(&uuid);
+                                continue;
+                            }
+
+                            // Transition to decode (no prefill cost)
+                            if sequence.generated_tokens() == 1 {
+                                state_guard.set_prefill_cost(uuid, None);
+                            }
+                        }
+
+                        // Sleep once for the accumulated duration
+                        if total_sleep_duration.as_millis() > 0 {
+                            tokio::time::sleep(total_sleep_duration).await;
+                        }
+                    }
+                }
+            }
+        });
+
+        Self {
+            state,
+            kv_manager,
+            token_capacity,
+            watermark,
+            block_size,
+            chunk_size,
+            background_handle: Some(background_handle),
+            request_tx,
+            cancellation_token,
+        }
+    }
+
+    /// Add a new request to the waiting queue
+    pub async fn receive(&self, request: DirectRequest) {
+        let _ = self.request_tx.send(request).await;
+    }
+
+    /// Get the count of waiting requests
+    pub async fn waiting_count(&self) -> usize {
+        let state = self.state.lock().await;
+        state.waiting.len()
+    }
+
+    /// Get the count of running requests
+    pub async fn running_count(&self) -> usize {
+        let state = self.state.lock().await;
+        state.running.len()
+    }
+
+    /// Get the current capacity of the KvManager
+    pub async fn kv_usage_perc(&self) -> f64 {
+        let kv_manager = self.kv_manager.lock().await;
+        kv_manager.current_capacity_perc()
+    }
+
+    /// Returns forward pass metrics for monitoring purposes
+    pub async fn get_forward_pass_metrics(&self) -> ForwardPassMetrics {
+        let state = self.state.lock().await;
+        let kv_manager = self.kv_manager.lock().await;
+
+        // Get the active blocks and total capacity from KvManager
+        let active_blocks_count = kv_manager.active_blocks().len() as u64;
+        let total_capacity = kv_manager.max_capacity() as u64;
+
+        // Calculate GPU cache usage percentage
+        let gpu_cache_usage_perc = if total_capacity > 0 {
+            active_blocks_count as f32 / total_capacity as f32
+        } else {
+            0.0
+        };
+
+        ForwardPassMetrics {
+            request_active_slots: state.running.len() as u64,
+            request_total_slots: 420, // Dummy value as specified
+            kv_active_blocks: active_blocks_count,
+            kv_total_blocks: total_capacity,
+            num_requests_waiting: state.waiting.len() as u64,
+            gpu_cache_usage_perc,
+            gpu_prefix_cache_hit_rate: 0.0, // Placeholder value as specified
+        }
+    }
+}
+
+// Implement Clone for Scheduler to support sharing between tasks
+impl Clone for Scheduler {
+    fn clone(&self) -> Self {
+        Self {
+            state: self.state.clone(),
+            kv_manager: self.kv_manager.clone(),
+            token_capacity: self.token_capacity,
+            watermark: self.watermark,
+            block_size: self.block_size,
+            chunk_size: self.chunk_size,
+            background_handle: None,
path: lib/llm/src/mocker/scheduler.rs
user: login: PeaBrane
body: I opted in the end to just get rid of the task handle to avoid bloat, assuming that the user would correctly pass in the `cancelation_token` as expected by the repo apparently
created_at: 2025-05-19T22:49:56Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,610 @@
+// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use crate::kv_router::protocols::ForwardPassMetrics;
+use crate::mocker::evictor::LRUEvictor;
+use crate::mocker::kv_manager::KvManager;
+use crate::mocker::protocols::DirectRequest;
+use crate::mocker::protocols::{MoveBlock, MoveBlockResponse, PrefillCost, UniqueBlock};
+use crate::mocker::sequence::ActiveSequence;
+use std::collections::HashMap;
+use std::collections::VecDeque;
+use std::sync::Arc;
+use tokio::sync::{mpsc, Mutex};
+use tokio::task::JoinHandle;
+use tokio::time::{interval, Duration};
+use tokio_util::sync::CancellationToken;
+use uuid::Uuid;
+
+/// Enum representing either a direct request or an active sequence
+pub enum Request {
+    Direct(DirectRequest),
+    Active(ActiveSequence),
+}
+struct SchedulerState {
+    waiting: VecDeque<Uuid>,
+    ready: VecDeque<Uuid>,
+    running: LRUEvictor<Uuid>,
+    requests: HashMap<Uuid, Request>,
+    prefill_costs: HashMap<Uuid, Option<PrefillCost>>,
+}
+
+impl Default for SchedulerState {
+    fn default() -> Self {
+        Self {
+            waiting: VecDeque::new(),
+            ready: VecDeque::new(),
+            running: LRUEvictor::new(),
+            requests: HashMap::new(),
+            prefill_costs: HashMap::new(),
+        }
+    }
+}
+
+impl SchedulerState {
+    /// Create a new UUID for a DirectRequest, add it to requests, and push the UUID to waiting.
+    fn receive(&mut self, request: DirectRequest) -> Uuid {
+        // Use the provided UUID if available, otherwise generate a new one
+        let uuid = request.uuid.unwrap_or_else(Uuid::new_v4);
+
+        // Add the request to the map and waiting queue
+        self.requests.insert(uuid, Request::Direct(request));
+        self.waiting.push_back(uuid);
+        uuid
+    }
+
+    /// Get the next UUID from ready or waiting queue and its associated Request.
+    /// Returns from ready if not empty, otherwise from waiting, or None if both are empty.
+    /// Also removes the Request from the requests HashMap.
+    fn next(&mut self) -> Option<(Uuid, Request)> {
+        let uuid = self
+            .ready
+            .pop_front()
+            .or_else(|| self.waiting.pop_front())?;
+        let request = self.requests.remove(&uuid)?;
+        Some((uuid, request))
+    }
+
+    /// Move a UUID and its Request to the ready queue.
+    fn make_ready(&mut self, uuid: Uuid, active_seq: ActiveSequence) {
+        self.requests.insert(uuid, Request::Active(active_seq));
+        self.ready.push_back(uuid);
+    }
+
+    /// Schedule the request with the given UUID.
+    /// Returns the creation signal from the ActiveSequence.
+    fn run(&mut self, uuid: Uuid, active_seq: ActiveSequence) -> MoveBlock {
+        // Insert the request into the map
+        self.requests.insert(uuid, Request::Active(active_seq));
+
+        // Get the creation signal
+        let Some(Request::Active(sequence)) = self.requests.get(&uuid) else {
+            panic!("Failed to get ActiveSequence for UUID");
+        };
+        let Some(signal) = sequence.creation_signal() else {
+            panic!("Failed to get creation signal from ActiveSequence");
+        };
+
+        // Add to running requests
+        self.running.insert(uuid);
+        signal.clone()
+    }
+
+    /// Set the prefill cost for a UUID
+    fn set_prefill_cost(&mut self, uuid: Uuid, cost: Option<PrefillCost>) {
+        self.prefill_costs.insert(uuid, cost);
+    }
+
+    /// Get the prefill compute value for a UUID if available
+    fn get_prefill_compute(&self, uuid: &Uuid) -> Option<f64> {
+        self.prefill_costs
+            .get(uuid)
+            .and_then(|cost| cost.as_ref())
+            .map(|cost| cost.prefill_compute)
+    }
+
+    /// Calculate the current running batched tokens
+    fn num_batched_tokens(&self) -> usize {
+        self.prefill_costs
+            .values()
+            .map(|cost| match cost {
+                Some(cost) => cost.new_tokens,
+                None => 1,
+            })
+            .sum()
+    }
+
+    /// Remove a UUID and its associated Request from collections.
+    fn complete(&mut self, uuid: &Uuid) {
+        self.running.remove(uuid);
+        self.requests.remove(uuid);
+        self.prefill_costs.remove(uuid);
+    }
+
+    /// Preempt the oldest running request by evicting it from running, resetting the sequence,
+    /// and adding it back to the waiting queue.
+    /// Returns the signal from reset_with_signal or None if no requests are running.
+    fn preempt(&mut self) -> Option<Vec<MoveBlock>> {
+        // Evict the oldest UUID from running
+        let uuid = self.running.evict()?;
+        eprintln!("Request {} will be preempted", uuid);
+
+        // Remove the request from the requests HashMap and ensure it's an ActiveSequence
+        let request = self.requests.remove(&uuid)?;
+
+        // Remove the prefill cost to force recomputation
+        self.prefill_costs.remove(&uuid);
+
+        // Extract the ActiveSequence from the Request enum
+        let Request::Active(active_sequence) = request else {
+            panic!("Expected ActiveSequence in running queue")
+        };
+
+        // Reset the sequence and get the new sequence and signal
+        let (new_sequence, signals) = active_sequence.reset_with_signal();
+
+        // Insert the new sequence back into the requests map and add to waiting queue
+        self.requests.insert(uuid, Request::Active(new_sequence));
+        self.waiting.push_back(uuid);
+
+        Some(signals)
+    }
+}
+
+/// Manages scheduling of requests using KvManager resources
+pub struct Scheduler {
+    state: Arc<Mutex<SchedulerState>>,
+    kv_manager: Arc<Mutex<KvManager>>,
+    token_capacity: usize,
+    watermark: f64,
+    block_size: usize,
+    chunk_size: usize,
+    background_handle: Option<JoinHandle<()>>,
+    request_tx: mpsc::Sender<DirectRequest>,
+    cancellation_token: CancellationToken,
+}
+
+impl Scheduler {
+    /// Create a new Scheduler with the given parameters
+    pub fn new(
+        kv_capacity: usize,
+        watermark: f64,
+        block_size: usize,
+        chunk_size: Option<usize>,
+        output_tx: Option<mpsc::Sender<Uuid>>,
+        cancellation_token: Option<CancellationToken>,
+    ) -> Self {
+        // Create KvManager internally
+        let kv_manager = KvManager::new(kv_capacity, block_size);
+
+        let token_capacity: usize = 8192;
+        let state = Arc::new(Mutex::new(SchedulerState::default()));
+
+        let kv_manager = Arc::new(Mutex::new(kv_manager));
+        let chunk_size = chunk_size.unwrap_or(256);
+
+        // Create channel for request handling
+        let (request_tx, mut request_rx) = mpsc::channel::<DirectRequest>(1024);
+
+        // Use provided cancellation token or create new one
+        let cancellation_token = cancellation_token.unwrap_or_default();
+        let token_clone = cancellation_token.clone();
+
+        // Create a clone for the background task
+        let state_clone = state.clone();
+        let kv_manager_clone = kv_manager.clone();
+        let output_tx_clone = output_tx.clone();
+
+        // Spawn main background task with cancellation token
+        let background_handle = tokio::spawn(async move {
path: lib/llm/src/mocker/scheduler.rs
user: login: PeaBrane
body: To debloat, opted to delegate the explicit runtime initialization logic to the Python bindings (in alignment with the current stack). In Rust, tokio would have its own runtime. 
created_at: 2025-05-19T22:52:05Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,378 @@
+// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use crate::mocker::evictor::LRUEvictor;
+use crate::mocker::protocols::{MoveBlock, PrefillCost, UniqueBlock};
+use crate::mocker::sequence::ActiveSequence;
+use derive_getters::Getters;
+use std::collections::{HashMap, HashSet};
+
+/// Mock implementation of worker for testing and simulation
+#[derive(Getters)]
+pub struct KvManager {
+    #[getter(copy)]
+    max_capacity: usize,
+
+    #[getter(copy)]
+    block_size: usize,
+
+    active_blocks: HashMap<UniqueBlock, usize>,
+
+    inactive_blocks: LRUEvictor<UniqueBlock>,
+
+    all_blocks: HashSet<UniqueBlock>,
+}
+
+impl KvManager {
+    pub fn new(max_capacity: usize, block_size: usize) -> Self {
+        let active_blocks = HashMap::new();
+        let inactive_blocks = LRUEvictor::default();
+        let all_blocks = HashSet::new();
+
+        KvManager {
+            max_capacity,
+            block_size,
+            active_blocks,
+            inactive_blocks,
+            all_blocks,
+        }
+    }
+
+    /// Process a MoveBlock instruction synchronously
+    pub fn process(&mut self, event: &MoveBlock) -> bool {
+        match event {
+            MoveBlock::Use(hashes, _) => {
+                for hash in hashes {
+                    // First check if it already exists in active blocks
+                    if let Some(ref_count) = self.active_blocks.get_mut(hash) {
+                        // Block already active, just increment reference count
+                        *ref_count += 1;
+                        continue;
+                    }
+
+                    // Then check if it exists in inactive and move it to active if found
+                    if self.inactive_blocks.remove(hash) {
+                        // Insert into active with reference count 1
+                        self.active_blocks.insert(hash.clone(), 1);
+                        continue;
+                    }
+
+                    // Get counts for capacity check
+                    let active_count = self.active_blocks.len();
+                    let inactive_count = self.inactive_blocks.len();
+
+                    // If at max capacity, evict the oldest entry from inactive blocks
+                    if active_count + inactive_count >= self.max_capacity {
+                        if let Some(evicted) = self.inactive_blocks.evict() {
+                            // Remove evicted block from all_blocks
+                            self.all_blocks.remove(&evicted);
+                        } else {
+                            // Return false instead of panicking
+                            return false;
+                        }
+                    }
+
+                    // Now insert the new block in active blocks with reference count 1
+                    self.active_blocks.insert(hash.clone(), 1);
+                    // Add to all_blocks as it's a new block
+                    self.all_blocks.insert(hash.clone());
+                }
+            }
+            MoveBlock::Destroy(hashes) => {
+                // Loop in inverse direction
+                for hash in hashes.iter().rev() {
+                    self.active_blocks.remove(hash).unwrap();
+                    // Remove from all_blocks when destroyed
+                    assert!(self.all_blocks.remove(hash));
+                }
+            }
+            MoveBlock::Deref(hashes) => {
+                // Loop in inverse direction
+                for hash in hashes.iter().rev() {
+                    // Decrement reference count and check if we need to move to inactive
+                    if let Some(ref_count) = self.active_blocks.get_mut(hash) {
+                        *ref_count -= 1;
+
+                        // If reference count reaches zero, remove from active and move to inactive
+                        if *ref_count == 0 {
+                            self.active_blocks.remove(hash);
+                            // Use the LRUEvictor's timing functionality
+                            self.inactive_blocks.insert(hash.clone());
+                        }
+                    }
+                }
+            }
+            MoveBlock::Promote(uuid, hash) => {
+                let uuid_block = UniqueBlock::PartialBlock(*uuid);
+                let hash_block = UniqueBlock::FullBlock(*hash);
+
+                let Some(ref_count) = self.active_blocks.remove(&uuid_block) else {
+                    panic!("Cannot promote non-existent block.");
+                };
+
+                // Replace with hash block, keeping the same reference count
+                self.active_blocks.insert(hash_block.clone(), ref_count);
+
+                // Update all_blocks
+                assert!(self.all_blocks.remove(&uuid_block));
+                self.all_blocks.insert(hash_block);
+            }
+        }
+
+        // Return true if we made it this far
+        true
+    }
+
+    /// Get the count of blocks in the input list that aren't in all_blocks
+    pub fn probe_new_blocks(&self, blocks: &[UniqueBlock]) -> usize {
+        blocks
+            .iter()
+            .filter(|&block| !self.all_blocks.contains(block))
+            .count()
+    }
+
+    /// Get the current capacity (active blocks + inactive blocks)
+    pub fn current_capacity(&self) -> usize {
+        let active = self.active_blocks.len();
+        let inactive = self.inactive_blocks.len();
+        active + inactive
+    }
+
+    /// Get the current capacity as a percentage of the maximum capacity
+    pub fn current_capacity_perc(&self) -> f64 {
+        let current = self.current_capacity() as f64;
+        current / self.max_capacity as f64
+    }
+
+    /// Get the keys of inactive blocks
+    pub fn get_inactive_blocks(&self) -> Vec<&UniqueBlock> {
+        self.inactive_blocks.keys().collect()
+    }
+
+    /// Get the keys of active blocks
+    pub fn get_active_blocks(&self) -> Vec<&UniqueBlock> {
+        self.active_blocks.keys().collect()
+    }
+
+    /// Check if a sequence can be scheduled and calculate cost if possible
+    pub fn try_schedule(
+        &self,
+        sequence: &ActiveSequence,
+        watermark: f64,
+        tokens_budget: usize,
+    ) -> Option<PrefillCost> {
+        // Return None immediately if tokens_budget is 0
+        if tokens_budget == 0 {
+            return None;
+        }
+
+        // Get unique blocks from the sequence
+        let unique_blocks = sequence.unique_blocks();
+
+        // Get the count of new blocks
+        let new_blocks = self.probe_new_blocks(unique_blocks);
+
+        // Calculate current usage and available capacity
+        let active_count = self.active_blocks.len();
+
+        // Check if we can schedule based on the watermark
+        if (active_count + new_blocks) as f64 > (1.0 - watermark) * self.max_capacity as f64 {
+            return None;
+        }
+
+        // Calculate overlap blocks
+        let overlap_blocks = unique_blocks.len() - new_blocks;
+
+        // Calculate new tokens
+        let new_tokens = sequence.num_input_tokens() - overlap_blocks * self.block_size;
+
+        // // Print the full equation with actual values substituted
+        // println!("{} = {} - ({} * {}) (new_tokens = num_input_tokens - overlap_blocks * block_size)",
+        //     new_tokens,
+        //     sequence.num_input_tokens(),
+        //     overlap_blocks,
+        //     self.block_size);
+
+        // Return None if new_tokens exceeds tokens_budget
+        if new_tokens > tokens_budget {
+            return None;
+        }
+
+        // Calculate prefill compute
+        let prefill_compute =
path: lib/llm/src/mocker/kv_manager.rs
user: login: alec-flowers
body: Can you add a note explaining the formula a bit more. This part seems key to the modeling, so worth explaining in some more detail. This number is also unitless, so I guess its something we can just compare against itself? 
created_at: 2025-05-21T00:12:23Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,378 @@
+// SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use crate::mocker::evictor::LRUEvictor;
+use crate::mocker::protocols::{MoveBlock, PrefillCost, UniqueBlock};
+use crate::mocker::sequence::ActiveSequence;
+use derive_getters::Getters;
+use std::collections::{HashMap, HashSet};
+
+/// Mock implementation of worker for testing and simulation
+#[derive(Getters)]
+pub struct KvManager {
+    #[getter(copy)]
+    max_capacity: usize,
+
+    #[getter(copy)]
+    block_size: usize,
+
+    active_blocks: HashMap<UniqueBlock, usize>,
+
+    inactive_blocks: LRUEvictor<UniqueBlock>,
+
+    all_blocks: HashSet<UniqueBlock>,
+}
+
+impl KvManager {
+    pub fn new(max_capacity: usize, block_size: usize) -> Self {
+        let active_blocks = HashMap::new();
+        let inactive_blocks = LRUEvictor::default();
+        let all_blocks = HashSet::new();
+
+        KvManager {
+            max_capacity,
+            block_size,
+            active_blocks,
+            inactive_blocks,
+            all_blocks,
+        }
+    }
+
+    /// Process a MoveBlock instruction synchronously
+    pub fn process(&mut self, event: &MoveBlock) -> bool {
+        match event {
+            MoveBlock::Use(hashes, _) => {
+                for hash in hashes {
+                    // First check if it already exists in active blocks
+                    if let Some(ref_count) = self.active_blocks.get_mut(hash) {
+                        // Block already active, just increment reference count
+                        *ref_count += 1;
+                        continue;
+                    }
+
+                    // Then check if it exists in inactive and move it to active if found
+                    if self.inactive_blocks.remove(hash) {
+                        // Insert into active with reference count 1
+                        self.active_blocks.insert(hash.clone(), 1);
+                        continue;
+                    }
+
+                    // Get counts for capacity check
+                    let active_count = self.active_blocks.len();
+                    let inactive_count = self.inactive_blocks.len();
+
+                    // If at max capacity, evict the oldest entry from inactive blocks
+                    if active_count + inactive_count >= self.max_capacity {
+                        if let Some(evicted) = self.inactive_blocks.evict() {
+                            // Remove evicted block from all_blocks
+                            self.all_blocks.remove(&evicted);
+                        } else {
+                            // Return false instead of panicking
+                            return false;
path: lib/llm/src/mocker/kv_manager.rs
user: login: alec-flowers
body: Add note that something needs to be evicted. I was initially confused by this comment. 
created_at: 2025-05-21T00:28:13Z
author_association: CONTRIBUTOR

PR comments:
user: login: copy-pr-bot[bot]
created_at: 2025-05-11T18:07:37Z
body: This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities [here](https://docs.gha-runners.nvidia.com/cpr/vetters).

Contributors can view more details about this message [here](https://docs.gha-runners.nvidia.com/cpr/contributors).

author_association: NONE
user: login: alec-flowers
created_at: 2025-05-21T00:34:09Z
body: We need to see how we can utilize this both for collecting numbers / building heuristics, and also in testing and Mock-ing things. 
author_association: CONTRIBUTOR
user: login: PeaBrane
created_at: 2025-05-21T00:43:02Z
body: > Can you add headers to each file describing the purpose and goal? You did a great job in the PR description. It would be good to translate that into the code.
> 
> Also I think would be useful once its set up to look at what sort of numbers can be generated by running these workers.
> 
> We may want to have it where the Block Manager actually emits Events that the KVRouter can receive the signal and do something with. This would mean adding functionality from the KVPublisher to the MockWorker.

Awesome. Yea, those would be targets / scopes for near-future PRs. I think I need to first hook this up to an `AsyncEngine`, then would try to test our current router impls with this mocker
author_association: CONTRIBUTOR