feat: Migrate to GDS MT backend (#1734)

Author: jthomson04

Cargo.lock

@@ -114,7 +114,7 @@ SGLANG_BASE_IMAGE_TAG="25.01-cuda12.8-devel-ubuntu24.04"
 VLLM_V1_BASE_IMAGE="nvcr.io/nvidia/cuda-dl-base"
 VLLM_V1_BASE_IMAGE_TAG="25.01-cuda12.8-devel-ubuntu24.04"
 
-NIXL_COMMIT=16348080f5bdeb9fe6058a23be140cec020ef3f3
+NIXL_COMMIT=fa800bcfe3814b08df9cda9c30443de8c19665e5
 NIXL_REPO=ai-dynamo/nixl.git
 
 NIXL_UCX_EFA_REF=7ec95b95e524a87e81cac92f5ca8523e3966b16b

container/build.sh

@@ -3,6 +3,7 @@
 
 use super::*;
 
+use std::sync::Arc;
 use utils::get_barrier_id;
 
 use llm_rs::block_manager::distributed::{KvbmWorker as KvbmWorkerImpl, KvbmWorkerConfig};
@@ -91,11 +92,11 @@ impl KvbmWorker {
         worker_id: usize,
         dtype_width_bytes: usize,
     ) -> PyResult<Self> {
-        let mut vllm_tensors: Vec<Box<dyn TorchTensor>> = Vec::with_capacity(tensors.len());
+        let mut vllm_tensors: Vec<Arc<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));
+            vllm_tensors.push(Arc::new(vllm_tensor));
         }
 
         let barrier_id = get_barrier_id();

lib/bindings/python/Cargo.lock

@@ -84,7 +84,7 @@ regex = "1"
 rayon = "1"
 
 # block_manager
-nixl-sys = {git="https://github.com/ai-dynamo/nixl", rev = "a7c654d46a14cd5ce635cc8c02433d71df93dedf", optional = true }
+nixl-sys = {git="https://github.com/ai-dynamo/nixl", rev = "fa800bcfe3814b08df9cda9c30443de8c19665e5", optional = true }
 cudarc = { version = "0.16.2", features = ["cuda-12020"], optional = true }
 ndarray = { version = "0.16", optional = true }
 nix = { version = "0.26", optional = true }

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

@@ -37,8 +37,7 @@ mod tests {
 
     use dynamo_runtime::logging::init as init_logging;
 
-    const NUM_DEVICE_BLOCKS: usize = 8;
-    const NUM_HOST_BLOCKS: usize = 8;
+    const NUM_BLOCKS: usize = 8;
 
     #[derive(Clone, Debug)]
     struct MockTensor {
@@ -100,12 +99,12 @@ mod tests {
         let barrier_id = get_unique_barrier_id();
 
         for i in 0..num_workers {
-            let tensors: Vec<Box<dyn TorchTensor>> =
-                vec![Box::new(MockTensor::new(vec![2, NUM_DEVICE_BLOCKS, 4096]))];
+            let tensors: Vec<Arc<dyn TorchTensor>> =
+                vec![Arc::new(MockTensor::new(vec![2, NUM_BLOCKS, 4096]))];
 
             let config = KvbmWorkerConfig::builder()
                 .barrier_id(barrier_id.clone())
-                .num_device_blocks(NUM_DEVICE_BLOCKS)
+                .num_device_blocks(NUM_BLOCKS)
                 .tensors(tensors)
                 .worker_id(i)
                 .build()?;
@@ -117,7 +116,8 @@ mod tests {
         let leader_config = KvbmLeaderConfig::builder()
             .barrier_id(barrier_id)
             .world_size(num_workers)
-            .num_host_blocks(NUM_HOST_BLOCKS)
+            .num_host_blocks(NUM_BLOCKS)
+            .num_disk_blocks(NUM_BLOCKS)
             .build()?;
 
         // When/if this returns, we know that all the workers were also successful.
@@ -137,7 +137,9 @@ mod tests {
 
         let (leader, _workers) = build_leader_and_workers(num_workers).await?;
 
-        for block_idx in 0..std::cmp::min(NUM_DEVICE_BLOCKS, NUM_HOST_BLOCKS) {
+        // Do a whole bunch of distributed transfers.
+
+        for block_idx in 0..NUM_BLOCKS {
             leader
                 .transfer_blocks_request(utils::BlockTransferRequest::new(
                     utils::BlockTransferPool::Device,
@@ -148,10 +150,21 @@ mod tests {
                 .await?;
         }
 
-        for block_idx in 0..std::cmp::min(NUM_DEVICE_BLOCKS, NUM_HOST_BLOCKS) {
+        for block_idx in 0..NUM_BLOCKS {
             leader
                 .transfer_blocks_request(utils::BlockTransferRequest::new(
                     utils::BlockTransferPool::Host,
+                    utils::BlockTransferPool::Disk,
+                    vec![(block_idx, block_idx)],
+                ))
+                .await?
+                .await?;
+        }
+
+        for block_idx in 0..NUM_BLOCKS {
+            leader
+                .transfer_blocks_request(utils::BlockTransferRequest::new(
+                    utils::BlockTransferPool::Disk,
                     utils::BlockTransferPool::Device,
                     vec![(block_idx, block_idx)],
                 ))
@@ -194,13 +207,19 @@ mod tests {
             )
             .device_layout(
                 KvManagerLayoutConfig::builder()
-                    .num_blocks(NUM_DEVICE_BLOCKS)
+                    .num_blocks(NUM_BLOCKS)
                     .logical(Some(BlockParallelismStrategy::LeaderWorkerSharded))
                     .build()?,
             )
             .host_layout(
                 KvManagerLayoutConfig::builder()
-                    .num_blocks(NUM_HOST_BLOCKS)
+                    .num_blocks(NUM_BLOCKS)
+                    .logical(Some(BlockParallelismStrategy::LeaderWorkerSharded))
+                    .build()?,
+            )
+            .disk_layout(
+                KvManagerLayoutConfig::builder()
+                    .num_blocks(NUM_BLOCKS)
                     .logical(Some(BlockParallelismStrategy::LeaderWorkerSharded))
                     .build()?,
             )
@@ -218,8 +237,9 @@ mod tests {
 
         let device_pool = block_manager.device().unwrap();
         let host_pool = block_manager.host().unwrap();
+        let disk_pool = block_manager.disk().unwrap();
 
-        let mut device_blocks = device_pool.allocate_blocks(NUM_DEVICE_BLOCKS).await?;
+        let mut device_blocks = device_pool.allocate_blocks(NUM_BLOCKS).await?;
 
         let mut sequence_hashes = Vec::new();
         for block in &mut device_blocks {
@@ -245,15 +265,21 @@ mod tests {
             .match_sequence_hashes(sequence_hashes.as_slice())
             .await?;
 
-        assert_eq!(host_blocks.len(), NUM_DEVICE_BLOCKS);
+        assert_eq!(host_blocks.len(), NUM_BLOCKS);
+
+        let disk_blocks = disk_pool
+            .match_sequence_hashes(sequence_hashes.as_slice())
+            .await?;
+
+        assert_eq!(disk_blocks.len(), NUM_BLOCKS);
 
         // Return the device blocks to the pool.
         drop(immutable_device_blocks);
 
         tokio::time::sleep(std::time::Duration::from_millis(100)).await;
 
         // Clear out the device pool.
-        let _ = device_pool.allocate_blocks(NUM_DEVICE_BLOCKS).await?;
+        let _ = device_pool.allocate_blocks(NUM_BLOCKS).await?;
 
         // Now, all the blocks should be gone.
         assert_eq!(
@@ -270,7 +296,7 @@ mod tests {
         // Now, onboard them back to the device.
         let new_device_blocks = block_manager.onboard_blocks(host_blocks, None).await??;
 
-        assert_eq!(new_device_blocks.len(), NUM_DEVICE_BLOCKS);
+        assert_eq!(new_device_blocks.len(), NUM_BLOCKS);
 
         Ok(())
     }

lib/llm/Cargo.toml

@@ -17,6 +17,8 @@ use std::time::Duration;
 use tokio::sync::oneshot;
 use tokio_util::sync::CancellationToken;
 
+const INIT_TIMEOUT_SECS: u64 = 120;
+
 /// Data that is sent to workers over ETCD to establish a ZMQ connection.
 #[derive(Debug, Clone, Serialize, Deserialize)]
 pub struct KvbmLeaderData {
@@ -86,7 +88,7 @@ impl KvbmLeader {
         let leader_barrier: LeaderBarrier<KvbmLeaderData, ()> = LeaderBarrier::new(
             config.barrier_id.clone(),
             config.world_size,
-            Some(Duration::from_secs(30)),
+            Some(Duration::from_secs(INIT_TIMEOUT_SECS)),
         );
 
         let worker_data = leader_barrier
@@ -103,7 +105,7 @@ impl KvbmLeader {
         let zmq_leader = ZmqActiveMessageLeader::new(
             leader_sockets,
             config.world_size,
-            Duration::from_secs(30),
+            Duration::from_secs(INIT_TIMEOUT_SECS),
             cancel_token.clone(),
         )
         .await?;

lib/llm/src/block_manager/distributed.rs

@@ -31,7 +31,7 @@ use dynamo_runtime::{
 };
 
 fn load_and_validate_tensors(
-    tensors: Vec<Box<dyn TorchTensor>>,
+    tensors: &[Arc<dyn TorchTensor>],
     device_id: usize,
 ) -> anyhow::Result<(Vec<DeviceStorage>, Vec<usize>)> {
     let mut shape = None;
@@ -67,7 +67,7 @@ fn load_and_validate_tensors(
         }
 
         // Build the storage object from the tensor.
-        let device_tensor = DeviceStorage::new_from_torch(allocator.ctx(), tensor)?;
+        let device_tensor = DeviceStorage::new_from_torch(allocator.ctx(), tensor.clone())?;
 
         device_tensors.push(device_tensor);
     }
@@ -84,7 +84,7 @@ pub struct KvbmWorkerConfig {
     page_size: usize,
 
     #[builder(default = "Vec::new()")]
-    tensors: Vec<Box<dyn TorchTensor>>,
+    tensors: Vec<Arc<dyn TorchTensor>>,
 
     #[builder(default = "0")]
     device_id: usize,
@@ -105,14 +105,13 @@ impl KvbmWorkerConfig {
     }
 }
 
-fn build_agent(worker_id: usize) -> anyhow::Result<NixlAgent> {
-    // TODO: Get GDS enabled here.
-    // There seems to be some issue with NIXL that causes errors if a large amount of GDS backends are instantiated all at once.
-
+fn build_agent(worker_id: usize, use_gds: bool) -> anyhow::Result<NixlAgent> {
     let agent = NixlAgent::new(&format!("kvbm-worker-{}", worker_id))?;
-    // let (_, gds_params) = agent.get_plugin_params("GDS")?;
+    if use_gds {
+        let (_, gds_params) = agent.get_plugin_params("GDS_MT")?;
+        agent.create_backend("GDS_MT", &gds_params)?;
+    }
     let (_, posix_params) = agent.get_plugin_params("POSIX")?;
-    // agent.create_backend("GDS", &gds_params)?;
     agent.create_backend("POSIX", &posix_params)?;
 
     Ok(agent)
@@ -135,7 +134,7 @@ impl KvbmWorker {
             return Err(anyhow::anyhow!("num_device_blocks must be greater than 0"));
         }
 
-        let (device_tensors, shape) = load_and_validate_tensors(config.tensors, config.device_id)?;
+        let (device_tensors, shape) = load_and_validate_tensors(&config.tensors, config.device_id)?;
 
         if shape.len() < 3 {
             return Err(anyhow::anyhow!(format!(
@@ -182,28 +181,14 @@ impl KvbmWorker {
 
         let layout_builder_clone = layout_builder.clone();
 
-        let agent = build_agent(config.worker_id)?;
-
-        let transfer_context = Arc::new(TransferContext::new(
-            Arc::new(Some(agent)),
-            DeviceAllocator::new(config.device_id)
-                .unwrap()
-                .ctx()
-                .new_stream()
-                .unwrap(),
-            Handle::current(),
-        ));
-
         let cancel_token = CancellationToken::new();
         let task = CriticalTaskExecutionHandle::new(
             move |cancel_token| {
                 KvbmWorker::worker_task(
                     device_layout,
                     layout_builder_clone,
                     layout_type,
-                    config.barrier_id,
-                    config.worker_id,
-                    transfer_context,
+                    config,
                     cancel_token,
                 )
             },
@@ -235,26 +220,22 @@ impl KvbmWorker {
         device_layout: Box<dyn NixlLayout<StorageType = DeviceStorage>>,
         mut layout_builder: LayoutConfigBuilder,
         layout_type: LayoutType,
-        barrier_id: String,
-        worker_id: usize,
-        transfer_context: Arc<TransferContext>,
+        config: KvbmWorkerConfig,
         cancel_token: CancellationToken,
     ) -> anyhow::Result<()> {
-        // Build our device, host, and disk block lists.
-        let device_blocks = Some(Self::make_layout::<_, BasicMetadata>(
-            device_layout,
-            transfer_context.nixl_agent().as_ref(),
-            0,
-            worker_id,
-        )?);
-
         let runtime = Runtime::from_current()?;
         let drt = DistributedRuntime::from_settings(runtime).await?;
 
-        tracing::info!("Worker {} waiting on barrier {}", worker_id, barrier_id);
+        tracing::info!(
+            "Worker {} waiting on barrier {}",
+            config.worker_id,
+            config.barrier_id
+        );
 
-        let worker_barrier =
-            WorkerBarrier::<KvbmLeaderData, ()>::new(barrier_id, worker_id.to_string());
+        let worker_barrier = WorkerBarrier::<KvbmLeaderData, ()>::new(
+            config.barrier_id,
+            config.worker_id.to_string(),
+        );
 
         let leader_data = tokio::select! {
             _ = cancel_token.cancelled() => {
@@ -268,10 +249,30 @@ impl KvbmWorker {
 
         tracing::info!(
             "Worker {} received leader data: {:?}",
-            worker_id,
+            config.worker_id,
             leader_data
         );
 
+        let agent = build_agent(config.worker_id, leader_data.num_disk_blocks > 0)?;
+
+        let transfer_context = Arc::new(TransferContext::new(
+            Arc::new(Some(agent)),
+            DeviceAllocator::new(config.device_id)
+                .unwrap()
+                .ctx()
+                .new_stream()
+                .unwrap(),
+            Handle::current(),
+        ));
+
+        // Build our device, host, and disk block lists.
+        let device_blocks = Some(Self::make_layout::<_, BasicMetadata>(
+            device_layout,
+            transfer_context.nixl_agent().as_ref(),
+            0,
+            config.worker_id,
+        )?);
+
         let host_blocks = if leader_data.num_host_blocks > 0 {
             let host_allocator = Arc::new(PinnedAllocator::default());
             let host_layout = layout_builder
@@ -283,7 +284,7 @@ impl KvbmWorker {
                 host_layout,
                 transfer_context.nixl_agent().as_ref(),
                 1,
-                worker_id,
+                config.worker_id,
             )?)
         } else {
             None
@@ -300,7 +301,7 @@ impl KvbmWorker {
                 disk_layout,
                 transfer_context.nixl_agent().as_ref(),
                 2,
-                worker_id,
+                config.worker_id,
             )?)
         } else {
             None