fix: handling 0 block transfer (#3613)

Signed-off-by: Olga Andreeva <oandreeva@nvidia.com>

lib/llm/src/block_manager/block/data/logical/distributed_leader_worker.rs

@@ -100,6 +100,20 @@ impl LogicalResources for DistributedLeaderWorkerResources {
         RB: BlockDataProvider<Locality = Logical<Self>>,
         WB: WritableBlock + BlockDataProviderMut<Locality = Logical<Self>>,
     {
+        // Check for empty slices and length mismatch early
+        if sources.is_empty() && targets.is_empty() {
+            tracing::warn!(
+                "DistributedLeaderWorkerResources::handle_transfer called with both sources and targets empty, skipping transfer"
+            );
+            let (tx, rx) = oneshot::channel();
+            tx.send(()).unwrap();
+            return Ok(rx);
+        }
+
+        if sources.len() != targets.len() {
+            return Err(TransferError::CountMismatch(sources.len(), targets.len()));
+        }
+
         if let Some(transfer_tx) = &self.transfer_tx {
             let source_pool = Self::get_pool(sources[0].block_data());
             let target_pool = Self::get_pool(targets[0].block_data());

lib/llm/src/block_manager/block/locality.rs

@@ -125,6 +125,20 @@ impl<R: LogicalResources> LocalityProvider for Logical<R> {
         RB: BlockDataProvider<Locality = Self>,
         WB: WritableBlock + BlockDataProviderMut<Locality = Self>,
     {
+        // Check for empty slices and length mismatch early
+        if sources.is_empty() && targets.is_empty() {
+            tracing::warn!(
+                "Logical::handle_transfer called with both sources and targets empty, skipping transfer"
+            );
+            let (tx, rx) = oneshot::channel();
+            tx.send(()).unwrap();
+            return Ok(rx);
+        }
+
+        if sources.len() != targets.len() {
+            return Err(TransferError::CountMismatch(sources.len(), targets.len()));
+        }
+
         let source_resources = Self::load_resources(sources);
         let target_resources = Self::load_resources_mut(targets);
 

lib/llm/src/block_manager/block/transfer.rs

@@ -179,6 +179,20 @@ where
     <RB as StorageTypeProvider>::StorageType: NixlDescriptor,
     <WB as StorageTypeProvider>::StorageType: NixlDescriptor,
 {
+    // Check for empty slices and length mismatch early
+    if sources.is_empty() && targets.is_empty() {
+        tracing::warn!(
+            "handle_local_transfer called with both sources and targets empty, skipping transfer"
+        );
+        let (tx, rx) = oneshot::channel();
+        tx.send(()).unwrap();
+        return Ok(rx);
+    }
+
+    if sources.len() != targets.len() {
+        return Err(TransferError::CountMismatch(sources.len(), targets.len()));
+    }
+
     let (tx, rx) = oneshot::channel();
 
     match RB::write_to_strategy() {