feats: add get_offloaded_computed_blocks to get reusable g2 g3 blocks (#1709)

Co-authored-by: jthomson04 <jwillthomson19@gmail.com>

Author: ziqifan617

lib/bindings/python/rust/llm/block_manager.rs

@@ -44,24 +44,6 @@ pub fn add_to_module(m: &Bound<'_, PyModule>) -> PyResult<()> {
     Ok(())
 }
 
-pub fn map_dtype(dtype: &str) -> anyhow::Result<dynamo_llm::common::dtype::DType> {
-    Ok(match dtype {
-        "fp8" | "FP8" => dynamo_llm::common::dtype::DType::FP8,
-        "fp16" | "FP16" => dynamo_llm::common::dtype::DType::FP16,
-        "bf16" | "BF16" => dynamo_llm::common::dtype::DType::BF16,
-        "fp32" | "FP32" => dynamo_llm::common::dtype::DType::FP32,
-        "u8" | "U8" => dynamo_llm::common::dtype::DType::U8,
-        "u16" | "U16" => dynamo_llm::common::dtype::DType::U16,
-        "u32" | "U32" => dynamo_llm::common::dtype::DType::U32,
-        "u64" | "U64" => dynamo_llm::common::dtype::DType::U64,
-        "i8" | "I8" => dynamo_llm::common::dtype::DType::I8,
-        "i16" | "I16" => dynamo_llm::common::dtype::DType::I16,
-        "i32" | "I32" => dynamo_llm::common::dtype::DType::I32,
-        "i64" | "I64" => dynamo_llm::common::dtype::DType::I64,
-        _ => return Err(anyhow::anyhow!("Unsupported dtype: {}", dtype)),
-    })
-}
-
 type VllmBlockManager = dynamo_llm::block_manager::KvBlockManager<
     Logical<DistributedLeaderWorkerResources>,
     BasicMetadata,
@@ -71,6 +53,7 @@ type VllmBlockManager = dynamo_llm::block_manager::KvBlockManager<
 #[derive(Clone)]
 pub struct BlockManager {
     inner: Arc<VllmBlockManager>,
+    _rt: Arc<tokio::runtime::Runtime>,
 }
 
 #[pymethods]
@@ -94,7 +77,7 @@ impl BlockManager {
 
         tracing::info!("Using {} device blocks", device_num_blocks);
 
-        let mut model_config = dynamo_llm::block_manager::KvManagerModelConfig::builder()
+        let model_config = dynamo_llm::block_manager::KvManagerModelConfig::builder()
             .num_layers(1)
             .outer_dim(1)
             .page_size(page_size)
@@ -110,47 +93,46 @@ impl BlockManager {
                 .map_err(to_pyerr)?,
         );
 
-        if leader.inner().num_host_blocks() > 0 {
-            tracing::info!("Using {} host blocks", leader.inner().num_host_blocks());
+        let (leader, rt) = leader.dissolve();
+
+        if leader.num_host_blocks() > 0 {
+            tracing::info!("Using {} host blocks", leader.num_host_blocks());
             config = config.host_layout(
                 dynamo_llm::block_manager::KvManagerLayoutConfig::builder()
-                    .num_blocks(leader.inner().num_host_blocks())
+                    .num_blocks(leader.num_host_blocks())
                     .logical(Some(BlockParallelismStrategy::LeaderWorkerSharded))
                     .build()
                     .map_err(to_pyerr)?,
             );
         }
 
-        if leader.inner().num_disk_blocks() > 0 {
-            tracing::info!("Using {} disk blocks", leader.inner().num_disk_blocks());
+        if leader.num_disk_blocks() > 0 {
+            tracing::info!("Using {} disk blocks", leader.num_disk_blocks());
             config = config.disk_layout(
                 dynamo_llm::block_manager::KvManagerLayoutConfig::builder()
-                    .num_blocks(leader.inner().num_disk_blocks())
+                    .num_blocks(leader.num_disk_blocks())
                     .logical(Some(BlockParallelismStrategy::LeaderWorkerSharded))
                     .build()
                     .map_err(to_pyerr)?,
             );
         }
 
         let config = config.build().map_err(to_pyerr)?;
-        let tokio_runtime = pyo3_async_runtimes::tokio::get_runtime();
         Ok(BlockManager {
             inner: Arc::from(
-                tokio_runtime
-                    .block_on(async {
-                        let resources = DistributedLeaderWorkerResources::new(
-                            leader.inner(),
-                            cancel_token.child_token(),
-                        )?;
-
-                        dynamo_llm::block_manager::KvBlockManager::<
-                            Logical<DistributedLeaderWorkerResources>,
-                            BasicMetadata,
-                        >::new(config, resources)
-                        .await
-                    })
-                    .map_err(to_pyerr)?,
+                rt.block_on(async {
+                    let resources =
+                        DistributedLeaderWorkerResources::new(leader, cancel_token.child_token())?;
+
+                    dynamo_llm::block_manager::KvBlockManager::<
+                        Logical<DistributedLeaderWorkerResources>,
+                        BasicMetadata,
+                    >::new(config, resources)
+                    .await
+                })
+                .map_err(to_pyerr)?,
             ),
+            _rt: rt,
         })
     }
 

lib/bindings/python/rust/llm/block_manager/distributed.rs

@@ -4,6 +4,7 @@
 use super::*;
 
 mod leader;
+mod utils;
 mod worker;
 
 pub use leader::KvbmLeader;

lib/bindings/python/rust/llm/block_manager/distributed/leader.rs

@@ -2,7 +2,9 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use super::*;
+use utils::get_barrier_id;
 
+use derive_getters::Dissolve;
 use llm_rs::block_manager::distributed::{KvbmLeader as KvbmLeaderImpl, KvbmLeaderConfig};
 
 fn compute_num_blocks(env_var: &str, bytes_per_block: usize) -> usize {
@@ -14,26 +16,22 @@ fn compute_num_blocks(env_var: &str, bytes_per_block: usize) -> usize {
 }
 
 #[pyclass]
-#[derive(Clone)]
+#[derive(Clone, Dissolve)]
 pub struct KvbmLeader {
     leader: Arc<KvbmLeaderImpl>,
-    _rt: Arc<tokio::runtime::Runtime>,
-}
-
-impl KvbmLeader {
-    pub fn inner(&self) -> Arc<KvbmLeaderImpl> {
-        self.leader.clone()
-    }
+    rt: Arc<tokio::runtime::Runtime>,
 }
 
 #[pymethods]
 impl KvbmLeader {
     #[new]
-    #[pyo3(signature = (barrier_id, bytes_per_block, world_size))]
-    fn new(barrier_id: String, bytes_per_block: usize, world_size: usize) -> PyResult<Self> {
+    #[pyo3(signature = (bytes_per_block, world_size))]
+    fn new(bytes_per_block: usize, world_size: usize) -> PyResult<Self> {
         let num_host_blocks = compute_num_blocks("DYNAMO_KVBM_CPU_CACHE", bytes_per_block);
         let num_disk_blocks = compute_num_blocks("DYNAMO_KVBM_DISK_CACHE", bytes_per_block);
 
+        let barrier_id = get_barrier_id();
+
         let config = KvbmLeaderConfig::builder()
             .barrier_id(barrier_id)
             .num_host_blocks(num_host_blocks)
@@ -52,7 +50,7 @@ impl KvbmLeader {
 
         Ok(Self {
             leader: Arc::new(leader),
-            _rt: Arc::new(rt),
+            rt: Arc::new(rt),
         })
     }
 }

lib/bindings/python/rust/llm/block_manager/distributed/utils.rs

@@ -0,0 +1,6 @@
+// SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+pub fn get_barrier_id() -> String {
+    std::env::var("DYNAMO_KVBM_BARRIER_ID").unwrap_or("kvbm".to_string())
+}

lib/bindings/python/rust/llm/block_manager/distributed/worker.rs

@@ -3,6 +3,8 @@
 
 use super::*;
 
+use utils::get_barrier_id;
+
 use llm_rs::block_manager::distributed::{KvbmWorker as KvbmWorkerImpl, KvbmWorkerConfig};
 use llm_rs::block_manager::storage::torch::{TorchDevice, TorchTensor};
 
@@ -80,32 +82,31 @@ pub struct KvbmWorker {
 #[pymethods]
 impl KvbmWorker {
     #[new]
-    #[pyo3(signature = (num_device_blocks, page_size, tensors, device_id=0, worker_id=0, dtype=None, barrier_id="kvbm".to_string()))]
+    #[pyo3(signature = (num_device_blocks, page_size, tensors, device_id=0, worker_id=0, dtype_width_bytes=2))]
     fn new(
         num_device_blocks: usize,
         page_size: usize,
         tensors: Vec<Py<PyAny>>,
         device_id: usize,
         worker_id: usize,
-        dtype: Option<String>,
-        barrier_id: String,
+        dtype_width_bytes: usize,
     ) -> PyResult<Self> {
-        let dtype = map_dtype(&dtype.unwrap_or("fp16".to_string())).map_err(to_pyerr)?;
-
         let mut vllm_tensors: Vec<Box<dyn TorchTensor>> = Vec::with_capacity(tensors.len());
 
         for tensor in tensors {
             let vllm_tensor = VllmTensor::new(tensor.clone()).map_err(to_pyerr)?;
             vllm_tensors.push(Box::new(vllm_tensor));
         }
 
+        let barrier_id = get_barrier_id();
+
         let config = KvbmWorkerConfig::builder()
             .num_device_blocks(num_device_blocks)
             .page_size(page_size)
             .tensors(vllm_tensors)
             .device_id(device_id)
             .worker_id(worker_id)
-            .dtype(dtype)
+            .dtype_width_bytes(dtype_width_bytes)
             .barrier_id(barrier_id)
             .build()
             .map_err(to_pyerr)?;