feat: predictive active blocks for routing without load metrics (#1731)

Signed-off-by: Yan Ru Pei <yanrpei@gmail.com>
Co-authored-by: Alec <35311602+alec-flowers@users.noreply.github.com>

lib/llm/src/mocker/protocols.rs

@@ -21,27 +21,12 @@ use crate::kv_router::protocols::{
     ExternalSequenceBlockHash, KvCacheEventData, KvCacheRemoveData, KvCacheStoreData,
     KvCacheStoredBlockData, LocalBlockHash,
 };
+use crate::tokens::blocks::UniqueBlock;
 
 pub type Token = u32;
 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
 /// For Use and Promote variants, parent hash is the second field
 #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]

lib/llm/src/mocker/scheduler.rs

@@ -43,11 +43,10 @@
 use crate::kv_router::protocols::{ForwardPassMetrics, KvCacheEventData, KvStats, WorkerStats};
 use crate::mocker::evictor::LRUEvictor;
 use crate::mocker::kv_manager::KvManager;
-use crate::mocker::protocols::{
-    block_response_to_kv_event, MoveBlock, OutputSignal, PrefillCost, UniqueBlock,
-};
+use crate::mocker::protocols::{block_response_to_kv_event, MoveBlock, OutputSignal, PrefillCost};
 use crate::mocker::protocols::{DirectRequest, MockEngineArgs, MoveBlockResponse};
 use crate::mocker::sequence::ActiveSequence;
+use crate::tokens::blocks::UniqueBlock;
 use std::collections::HashMap;
 use std::collections::VecDeque;
 use std::sync::Arc;

lib/llm/src/mocker/sequence.rs

@@ -13,7 +13,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-use crate::mocker::protocols::{MoveBlock, UniqueBlock};
+use crate::mocker::protocols::MoveBlock;
+use crate::tokens::blocks::UniqueBlock;
 use crate::tokens::{TokenBlockSequence, Tokens};
 use derive_getters::Getters;
 use rand::random;

lib/llm/src/tokens.rs

@@ -22,6 +22,8 @@ use derive_getters::Dissolve;
 use rayon::prelude::*;
 use std::ops::Range;
 
+pub mod blocks;
+
 /// A token is represented as a 32-bit unsigned integer.
 pub type Token = u32;
 
@@ -845,6 +847,18 @@ impl TokenBlockSequence {
 
         (result_blocks, current_block)
     }
+
+    pub fn from_slice(tokens: &[Token], block_size: u32, salt_hash: Option<SaltHash>) -> Self {
+        assert!(block_size > 0, "block_size must be greater than 0");
+        let salt_hash = salt_hash.unwrap_or(0);
+        let (blocks, current_block) = Self::split_tokens(tokens, block_size, salt_hash);
+
+        Self {
+            blocks,
+            current_block,
+            salt_hash,
+        }
+    }
 }
 
 #[cfg(test)]

lib/llm/src/tokens/blocks.rs

@@ -12,3 +12,24 @@
 // 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 GlobalHash = u64;
+
+/// Represents an active block beign built
+#[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())
+    }
+}