diff --git a/pom.xml b/pom.xml
index d9330461..389cd465 100755
--- a/pom.xml
+++ b/pom.xml
@@ -35,6 +35,7 @@ See file LICENSE for terms.
     <scala.version>2.12.12</scala.version>
     <scala.compat.version>2.12</scala.compat.version>
     <jucx.version>1.10.0-SNAPSHOT</jucx.version>
+    <cudf.version>0.16</cudf.version>
   </properties>
 
   <dependencies>
@@ -55,6 +56,12 @@ See file LICENSE for terms.
       <version>3.2.1</version>
       <scope>test</scope>
     </dependency>
+    <dependency>
+      <groupId>ai.rapids</groupId>
+      <artifactId>cudf</artifactId>
+      <version>${cudf.version}</version>
+      <scope>test</scope>
+    </dependency>
   </dependencies>
 
   <build>
@@ -137,6 +144,18 @@ See file LICENSE for terms.
           </execution>
         </executions>
       </plugin>
+      <plugin>
+        <groupId>org.apache.maven.plugins</groupId>
+        <artifactId>maven-jar-plugin</artifactId>
+        <version>3.2.0</version>
+        <executions>
+          <execution>
+            <goals>
+              <goal>test-jar</goal>
+            </goals>
+          </execution>
+        </executions>
+      </plugin>
     </plugins>
   </build>
 
diff --git a/src/main/scala/org/apache/spark/shuffle/ucx/UcxShuffleTransport.scala b/src/main/scala/org/apache/spark/shuffle/ucx/UcxShuffleTransport.scala
index 10aab2f2..8d150b9e 100755
--- a/src/main/scala/org/apache/spark/shuffle/ucx/UcxShuffleTransport.scala
+++ b/src/main/scala/org/apache/spark/shuffle/ucx/UcxShuffleTransport.scala
@@ -153,7 +153,7 @@ class UcxShuffleTransport(var ucxShuffleConf: UcxShuffleConf = null, val executo
   private[ucx] def handlePrefetchRequest(workerId: String, workerAddress: ByteBuffer,
                                          blockIds: Seq[BlockId]) {
 
-    logInfo(s"Prefetching blocks: ${blockIds.mkString(",")}")
+    logDebug(s"Prefetching blocks: ${blockIds.mkString(",")}")
     clientConnections.getOrElseUpdate(workerId,
       globalWorker.newEndpoint(new UcpEndpointParams().setUcpAddress(workerAddress))
     )
@@ -184,7 +184,7 @@ class UcxShuffleTransport(var ucxShuffleConf: UcxShuffleConf = null, val executo
     lock.lock()
     val blockMemory = block.getMemoryBlock
 
-    logInfo(s"Sending $blockId of size ${blockMemory.size} to tag: $tag")
+    logDebug(s"Sending $blockId of size ${blockMemory.size} to tag: $tag")
     ep.sendTaggedNonBlocking(blockMemory.address, blockMemory.size, tag, new UcxCallback {
       override def onSuccess(request: UcpRequest): Unit = {
         if (block.isInstanceOf[UcxPinnedBlock]) {
diff --git a/src/main/scala/org/apache/spark/shuffle/ucx/UcxWorkerWrapper.scala b/src/main/scala/org/apache/spark/shuffle/ucx/UcxWorkerWrapper.scala
index 2b55eb06..d4337e71 100755
--- a/src/main/scala/org/apache/spark/shuffle/ucx/UcxWorkerWrapper.scala
+++ b/src/main/scala/org/apache/spark/shuffle/ucx/UcxWorkerWrapper.scala
@@ -5,6 +5,7 @@
 package org.apache.spark.shuffle.ucx
 
 import java.io.{Closeable, ObjectOutputStream}
+import java.nio.BufferOverflowException
 import java.util.concurrent.ThreadLocalRandom
 
 import scala.collection.mutable
@@ -91,7 +92,7 @@ class UcxWorkerWrapper(worker: UcpWorker, transport: UcxShuffleTransport, ucxCon
       logInfo(s"Worker from thread ${Thread.currentThread().getName} connecting to $executorId")
       val endpointParams = new UcpEndpointParams()
         .setUcpAddress(workerAdresses.get(executorId))
-     worker.newEndpoint(endpointParams)
+      worker.newEndpoint(endpointParams)
     })
   }
 
@@ -115,16 +116,16 @@ class UcxWorkerWrapper(worker: UcpWorker, transport: UcxShuffleTransport, ucxCon
 
     ep.sendTaggedNonBlocking(mem.address, transport.ucxShuffleConf.rpcMessageSize,
       UcxRpcMessages.PREFETCH_TAG, new UcxCallback() {
-      override def onSuccess(request: UcpRequest): Unit = {
-        logTrace(s"Sent prefetch ${blockIds.length} blocks to $executorId")
-        memoryPool.put(mem)
-      }
-    })
+        override def onSuccess(request: UcpRequest): Unit = {
+          logTrace(s"Sent prefetch ${blockIds.length} blocks to $executorId")
+          memoryPool.put(mem)
+        }
+      })
   }
 
   private[ucx] def fetchBlocksByBlockIds(executorId: String, blockIds: Seq[BlockId],
-                            resultBuffer: Seq[MemoryBlock],
-                            callbacks: Seq[OperationCallback]): Seq[Request] = {
+                                         resultBuffer: Seq[MemoryBlock],
+                                         callbacks: Seq[OperationCallback]): Seq[Request] = {
     val ep = getConnection(executorId)
     val mem = memoryPool.get(transport.ucxShuffleConf.rpcMessageSize)
     val buffer = UcxUtils.getByteBufferView(mem.address,
@@ -136,13 +137,20 @@ class UcxWorkerWrapper(worker: UcpWorker, transport: UcxShuffleTransport, ucxCon
 
     Utils.tryWithResource(new ByteBufferBackedOutputStream(buffer)) { bos =>
       val out = new ObjectOutputStream(bos)
-      out.writeObject(message)
-      out.flush()
-      out.close()
+      try {
+        out.writeObject(message)
+        out.flush()
+        out.close()
+      } catch {
+        case _: BufferOverflowException =>
+          throw new UcxException(s"Prefetch blocks message size > " +
+            s"${transport.ucxShuffleConf.RPC_MESSAGE_SIZE.key}:${transport.ucxShuffleConf.rpcMessageSize}")
+        case ex: Exception => throw new UcxException(ex.getMessage)
+      }
     }
 
     val tag = ThreadLocalRandom.current().nextLong(Long.MinValue, 0)
-    logInfo(s"Sending message to $executorId to fetch ${blockIds.length} blocks on tag $tag")
+    logTrace(s"Sending message to $executorId to fetch ${blockIds.length} blocks on tag $tag")
     ep.sendTaggedNonBlocking(mem.address, transport.ucxShuffleConf.rpcMessageSize, tag,
       new UcxCallback() {
         override def onSuccess(request: UcpRequest): Unit = {
@@ -179,7 +187,7 @@ class UcxWorkerWrapper(worker: UcpWorker, transport: UcxShuffleTransport, ucxCon
   }
 
   private[ucx] def fetchBlockByBlockId(executorId: String, blockId: BlockId,
-                          resultBuffer: MemoryBlock, cb: OperationCallback): UcxRequest = {
+                                       resultBuffer: MemoryBlock, cb: OperationCallback): UcxRequest = {
     val stats = new UcxStats()
     val ep = getConnection(executorId)
     val mem = memoryPool.get(transport.ucxShuffleConf.rpcMessageSize)
@@ -206,28 +214,28 @@ class UcxWorkerWrapper(worker: UcpWorker, transport: UcxShuffleTransport, ucxCon
         override def onSuccess(request: UcpRequest): Unit = {
           memoryPool.put(mem)
         }
-    })
+      })
 
     val result = new UcxSuccessOperationResult(stats)
     val request = worker.recvTaggedNonBlocking(resultBuffer.address, resultBuffer.size,
       tag, -1L, new UcxCallback () {
 
-      override def onError(ucsStatus: Int, errorMsg: String): Unit = {
-        logError(s"Failed to receive blockId $blockId on tag: $tag, from executorId: $executorId " +
-          s" of size: ${resultBuffer.size}: $errorMsg")
-        if (cb != null ) {
-          cb.onComplete(new UcxFailureOperationResult(errorMsg))
+        override def onError(ucsStatus: Int, errorMsg: String): Unit = {
+          logError(s"Failed to receive blockId $blockId on tag: $tag, from executorId: $executorId " +
+            s" of size: ${resultBuffer.size}: $errorMsg")
+          if (cb != null ) {
+            cb.onComplete(new UcxFailureOperationResult(errorMsg))
+          }
         }
-      }
 
-      override def onSuccess(request: UcpRequest): Unit = {
-        stats.endTime = System.nanoTime()
-        stats.receiveSize = request.getRecvSize
-        if (cb != null) {
-          cb.onComplete(result)
+        override def onSuccess(request: UcpRequest): Unit = {
+          stats.endTime = System.nanoTime()
+          stats.receiveSize = request.getRecvSize
+          if (cb != null) {
+            cb.onComplete(result)
+          }
         }
-      }
-    })
+      })
     new UcxRequest(request, stats)
   }
 
diff --git a/src/main/scala/org/apache/spark/shuffle/ucx/perf/UcxShuffleTransportPerfTool.scala b/src/main/scala/org/apache/spark/shuffle/ucx/perf/UcxShuffleTransportPerfTool.scala
new file mode 100755
index 00000000..039b3219
--- /dev/null
+++ b/src/main/scala/org/apache/spark/shuffle/ucx/perf/UcxShuffleTransportPerfTool.scala
@@ -0,0 +1,224 @@
+/*
+* Copyright (C) Mellanox Technologies Ltd. 2020. ALL RIGHTS RESERVED.
+* See file LICENSE for terms.
+*/
+package org.apache.spark.shuffle.ucx.perf
+
+import java.net.{InetAddress, InetSocketAddress, ServerSocket, Socket}
+import java.nio.ByteBuffer
+import java.util.concurrent.atomic.{AtomicInteger, AtomicLong}
+import java.util.concurrent.{CountDownLatch, Executors, TimeUnit}
+
+import org.apache.commons.cli.{GnuParser, HelpFormatter, Options}
+import org.apache.spark.SparkConf
+import org.apache.spark.shuffle.ucx._
+import org.apache.spark.shuffle.ucx.memory.MemoryPool
+import org.apache.spark.util.Utils
+
+object UcxShuffleTransportPerfTool {
+  private val HELP_OPTION = "h"
+  private val ADDRESS_OPTION = "a"
+  private val NUM_BLOCKS_OPTION = "n"
+  private val SIZE_OPTION = "s"
+  private val PORT_OPTION = "p"
+  private val ITER_OPTION = "i"
+  private val MEMORY_TYPE_OPTION = "m"
+  private val NUM_THREADS_OPTION = "t"
+
+  private val ucxShuffleConf = new UcxShuffleConf(new SparkConf())
+  private val transport = new UcxShuffleTransport(ucxShuffleConf, "e")
+  private val workerAddress = transport.init()
+  private var memoryPool: MemoryPool = transport.memoryPool
+
+  case class TestBlockId(id: Int) extends BlockId
+
+  case class PerfOptions(remoteAddress: InetSocketAddress, numBlocks: Int, blockSize: Long,
+                         serverPort: Int, numIterations: Int, numThreads: Int)
+
+  private def initOptions(): Options = {
+    val options = new Options()
+    options.addOption(HELP_OPTION, "help", false,
+      "display help message")
+    options.addOption(ADDRESS_OPTION, "address", true,
+      "address of the remote host")
+    options.addOption(NUM_BLOCKS_OPTION, "num-blocks", true,
+      "number of blocks to transfer. Default: 1")
+    options.addOption(SIZE_OPTION, "block-size", true,
+      "size of block to transfer. Default: 4m")
+    options.addOption(PORT_OPTION, "server-port", true,
+      "server port. Default: 12345")
+    options.addOption(ITER_OPTION, "num-iterations", true,
+      "number of iterations. Default: 5")
+    options.addOption(NUM_THREADS_OPTION, "num-threads", true,
+      "number of threads. Default: 1")
+    options.addOption(MEMORY_TYPE_OPTION, "memory-type", true,
+      "memory type: host (default), cuda")
+  }
+
+  private def parseOptions(args: Array[String]): PerfOptions = {
+    val parser = new GnuParser()
+    val options = initOptions()
+    val cmd = parser.parse(options, args)
+
+    if (cmd.hasOption(HELP_OPTION)) {
+      new HelpFormatter().printHelp("UcxShufflePerfTool", options)
+      System.exit(0)
+    }
+
+    val inetAddress = if (cmd.hasOption(ADDRESS_OPTION)) {
+      val Array(host, port) = cmd.getOptionValue(ADDRESS_OPTION).split(":")
+      new InetSocketAddress(host, Integer.parseInt(port))
+    } else {
+      null
+    }
+
+    val serverPort = Integer.parseInt(cmd.getOptionValue(PORT_OPTION, "12345"))
+
+    val numIterations =  Integer.parseInt(cmd.getOptionValue(ITER_OPTION, "5"))
+
+    val threadsNumber = Integer.parseInt(cmd.getOptionValue(NUM_THREADS_OPTION, "1"))
+
+    if (cmd.hasOption(MEMORY_TYPE_OPTION) && cmd.getOptionValue(MEMORY_TYPE_OPTION) == "cuda") {
+      val className = "org.apache.spark.shuffle.ucx.GpuMemoryPool"
+      val cls = Utils.classForName(className)
+      memoryPool = cls.getConstructor().newInstance().asInstanceOf[MemoryPool]
+    }
+
+    PerfOptions(inetAddress,
+      Integer.parseInt(cmd.getOptionValue(NUM_BLOCKS_OPTION, "1")),
+      Utils.byteStringAsBytes(cmd.getOptionValue(SIZE_OPTION, "4m")),
+      serverPort, numIterations, threadsNumber)
+  }
+
+  private def startServer(perfOptions: PerfOptions): Unit = {
+    val blocks: Seq[Block] = (0 until perfOptions.numBlocks).map { _ =>
+      val block = memoryPool.get(perfOptions.blockSize)
+      new Block {
+        override def getMemoryBlock: MemoryBlock =
+          block
+      }
+    }
+
+    val blockIds = (0 until perfOptions.numBlocks).map(i => TestBlockId(i))
+    blockIds.zip(blocks).foreach {
+      case (blockId, block) => transport.register(blockId, block)
+    }
+
+    val serverSocket = new ServerSocket(perfOptions.serverPort)
+
+    println(s"Waiting for connections on " +
+      s"${InetAddress.getLocalHost.getHostName}:${perfOptions.serverPort} ")
+
+    val clientSocket = serverSocket.accept()
+    val out = clientSocket.getOutputStream
+    val in = clientSocket.getInputStream
+
+    val buf = ByteBuffer.allocate(workerAddress.capacity())
+    buf.put(workerAddress)
+    buf.flip()
+
+    out.write(buf.array())
+    out.flush()
+
+    println(s"Sending worker address to ${clientSocket.getInetAddress}")
+
+    buf.flip()
+
+    in.read(buf.array())
+    clientSocket.close()
+    serverSocket.close()
+
+    blocks.foreach(block => memoryPool.put(block.getMemoryBlock))
+    blockIds.foreach(transport.unregister)
+    transport.close()
+  }
+
+  private def startClient(perfOptions: PerfOptions): Unit = {
+    val socket = new Socket(perfOptions.remoteAddress.getHostName,
+      perfOptions.remoteAddress.getPort)
+
+    val buf = new Array[Byte](4096)
+    val readSize = socket.getInputStream.read(buf)
+    val executorId = "1"
+    val workerAddress = ByteBuffer.allocateDirect(readSize)
+
+    workerAddress.put(buf, 0, readSize)
+    println("Received worker address")
+
+    transport.addExecutor(executorId, workerAddress)
+
+    val resultSize = perfOptions.numBlocks * perfOptions.blockSize
+    val resultMemory = memoryPool.get(resultSize)
+
+    val blockIds = (0 until perfOptions.numBlocks).map(i => TestBlockId(i))
+
+    val threadPool =  Executors.newFixedThreadPool(perfOptions.numThreads)
+
+    for (i <- 1 to perfOptions.numIterations) {
+      val elapsedTime = new AtomicLong(0)
+      val countDownLatch = new CountDownLatch(perfOptions.numThreads)
+
+      for (_ <- 0 until perfOptions.numThreads) {
+        threadPool.execute(() => {
+          val completed = new AtomicInteger(0)
+
+          val mem = new Array[MemoryBlock](perfOptions.numBlocks)
+          val callbacks = new Array[OperationCallback](perfOptions.numBlocks)
+
+          for (j <- 0 until perfOptions.numBlocks) {
+            mem(j) = MemoryBlock(resultMemory.address + j * perfOptions.blockSize, perfOptions.blockSize)
+            callbacks(j) = (result: OperationResult) => {
+              elapsedTime.addAndGet(result.getStats.get.getElapsedTimeNs)
+              completed.incrementAndGet()
+            }
+          }
+
+          transport.fetchBlocksByBlockIds(executorId, blockIds, mem, callbacks)
+
+          while (completed.get() != perfOptions.numBlocks) {
+            transport.progress()
+          }
+          countDownLatch.countDown()
+        })
+      }
+
+      countDownLatch.await()
+
+      val totalTime  = if (elapsedTime.get() < TimeUnit.MILLISECONDS.toNanos(1)) {
+        s"$elapsedTime ns"
+      } else {
+        s"${TimeUnit.NANOSECONDS.toMillis(elapsedTime.get())} ms"
+      }
+      val throughput: Double = (resultSize * perfOptions.numThreads / 1024.0D / 1024.0D / 1024.0D) /
+        (elapsedTime.get() / 1e9D)
+
+      println(f"${s"[$i/${perfOptions.numIterations}]"}%12s" +
+        s" numBlocks: ${perfOptions.numBlocks}" +
+        s" numThreads: ${perfOptions.numThreads}" +
+        s" size: ${Utils.bytesToString(perfOptions.blockSize)}," +
+        s" total size: ${Utils.bytesToString(resultSize * perfOptions.numThreads)}," +
+        f" time: $totalTime%3s" +
+        f" throughput: $throughput%.5f GB/s")
+    }
+
+    val out = socket.getOutputStream
+    out.write(buf)
+    out.flush()
+    out.close()
+    socket.close()
+
+    memoryPool.put(resultMemory)
+    transport.close()
+  }
+
+  def main(args: Array[String]): Unit = {
+    val perfOptions = parseOptions(args)
+
+    if (perfOptions.remoteAddress == null) {
+      startServer(perfOptions)
+    } else {
+      startClient(perfOptions)
+    }
+  }
+
+}
diff --git a/src/main/scala/org/apache/spark/shuffle/ucx/rpc/GlobalWorkerRpcThread.scala b/src/main/scala/org/apache/spark/shuffle/ucx/rpc/GlobalWorkerRpcThread.scala
index c65dba23..efff6d76 100755
--- a/src/main/scala/org/apache/spark/shuffle/ucx/rpc/GlobalWorkerRpcThread.scala
+++ b/src/main/scala/org/apache/spark/shuffle/ucx/rpc/GlobalWorkerRpcThread.scala
@@ -21,7 +21,7 @@ class GlobalWorkerRpcThread(globalWorker: UcpWorker, memPool: MemoryPool,
   setDaemon(true)
   setName("Ucx Shuffle Transport Progress Thread")
 
-
+  
   override def run(): Unit = {
     val numRecvs = transport.ucxShuffleConf.recvQueueSize
     val msgSize = transport.ucxShuffleConf.rpcMessageSize
diff --git a/src/test/scala/org/apache/spark/shuffle/ucx/GpuMemoryPool.scala b/src/test/scala/org/apache/spark/shuffle/ucx/GpuMemoryPool.scala
new file mode 100755
index 00000000..6d3e633a
--- /dev/null
+++ b/src/test/scala/org/apache/spark/shuffle/ucx/GpuMemoryPool.scala
@@ -0,0 +1,30 @@
+/*
+* Copyright (C) Mellanox Technologies Ltd. 2020. ALL RIGHTS RESERVED.
+* See file LICENSE for terms.
+*/
+package org.apache.spark.shuffle.ucx
+
+import ai.rapids.cudf.DeviceMemoryBuffer
+import org.apache.spark.shuffle.ucx.memory.MemoryPool
+
+
+/**
+ * Test GPU mempool to run with [[ UcxShuffleTransportPerfTool ]]
+ */
+class GpuMemoryPool extends MemoryPool {
+
+  class GpuMemoryBlock(val deviceBuffer: DeviceMemoryBuffer,
+                       override val address: Long, override val size: Long)
+    extends MemoryBlock(address, size, isHostMemory = false)
+
+  override def get(size: Long): MemoryBlock = {
+    val deviceBuffer = DeviceMemoryBuffer.allocate(size)
+    new GpuMemoryBlock(deviceBuffer, deviceBuffer.getAddress, size)
+  }
+
+  override def put(mem: MemoryBlock): Unit = {
+    mem.asInstanceOf[GpuMemoryBlock].deviceBuffer.close()
+  }
+
+  override def close(): Unit = ???
+}
diff --git a/src/test/scala/org/apache/spark/shuffle/ucx/UcxShuffleTransportTestSuite.scala b/src/test/scala/org/apache/spark/shuffle/ucx/UcxShuffleTransportTestSuite.scala
index f391f3cd..7997042a 100755
--- a/src/test/scala/org/apache/spark/shuffle/ucx/UcxShuffleTransportTestSuite.scala
+++ b/src/test/scala/org/apache/spark/shuffle/ucx/UcxShuffleTransportTestSuite.scala
@@ -1,5 +1,5 @@
 /*
-* Copyright (C) Mellanox Technologies Ltd. 2019. ALL RIGHTS RESERVED.
+* Copyright (C) Mellanox Technologies Ltd. 2020. ALL RIGHTS RESERVED.
 * See file LICENSE for terms.
 */
 package org.apache.spark.shuffle.ucx