From 19aadb074c47c4a6867d649b1abae83b53a0c388 Mon Sep 17 00:00:00 2001
From: Thomas Vegas <tvegas@nvidia.com>
Date: Mon, 10 Aug 2026 16:46:35 +0200
Subject: [PATCH] UCT/CUDA_IFACE: Add sm_107 and bandwidth value

---
 config/m4/cuda.m4                      |  6 +++++-
 config/ucx.conf                        |  2 ++
 src/uct/cuda/cuda_ipc/cuda_ipc_iface.c | 16 +++++++++++++++-
 3 files changed, 22 insertions(+), 2 deletions(-)

diff --git a/config/m4/cuda.m4 b/config/m4/cuda.m4
index 54cb103253c..9c062103374 100644
--- a/config/m4/cuda.m4
+++ b/config/m4/cuda.m4
@@ -34,6 +34,7 @@ AC_DEFUN([UCX_CUDA_CHECK_NVCC], [
     ARCH124_CODE="-gencode=arch=compute_89,code=sm_89"
     ARCH128_CODE="-gencode=arch=compute_100,code=sm_100 -gencode=arch=compute_120,code=sm_120"
     ARCH130_CODE="-gencode=arch=compute_110,code=sm_110"
+    ARCH134_CODE="-gencode=arch=compute_107,code=sm_107"
 
 
     ARCH9_PTX="-gencode=arch=compute_70,code=compute_70"
@@ -44,6 +45,7 @@ AC_DEFUN([UCX_CUDA_CHECK_NVCC], [
     ARCH124_PTX="-gencode=arch=compute_90,code=compute_90"
     ARCH128_PTX="-gencode=arch=compute_120,code=compute_120"
     ARCH130_PTX="-gencode=arch=compute_120,code=compute_120"
+    ARCH134_PTX="-gencode=arch=compute_107,code=compute_107"
 
     AS_IF([test "x$NVCC" != "x"], [
         CUDA_VERSION=$($NVCC --version | grep release | sed 's/.*release //' | sed 's/\,.*//')
@@ -75,7 +77,9 @@ AC_DEFUN([UCX_CUDA_CHECK_NVCC], [
                         [AS_CASE([$CUDA_MAJOR_VERSION],
                                  [13],
                                      [# offline compilation support for architectures before '<compute/sm/lto>_75' is discontinued
-                                      NVCC_ARCH="${ARCH10_CODE} ${ARCH110_CODE} ${ARCH111_CODE} ${ARCH120_CODE} ${ARCH124_CODE} ${ARCH128_CODE} ${ARCH130_CODE} ${ARCH130_PTX}"],
+                                      NVCC_ARCH="${ARCH10_CODE} ${ARCH110_CODE} ${ARCH111_CODE} ${ARCH120_CODE} ${ARCH124_CODE} ${ARCH128_CODE} ${ARCH130_CODE} ${ARCH130_PTX}"
+                                      AS_CASE([$CUDA_MINOR_VERSION],
+                                              [4], [NVCC_ARCH="${NVCC_ARCH} ${ARCH134_CODE} ${ARCH134_PTX}"])],
                                  [12],
                                      [AS_CASE([$CUDA_MINOR_VERSION],
                                               [0|1|2|3],
diff --git a/config/ucx.conf b/config/ucx.conf
index ae91ee4a1a6..754b2c40290 100644
--- a/config/ucx.conf
+++ b/config/ucx.conf
@@ -20,6 +20,8 @@ UCX_MAX_RMA_RAILS=4
 UCX_REG_NONBLOCK_MEM_TYPES=host,cuda-managed
 UCX_IB_ODP_MEM_TYPES=host,cuda-managed
 UCX_IB_PCI_RELAXED_ORDERING=try
+UCX_CUDA_COPY_ASYNC_MEM_TYPE=cuda
+UCX_CUDA_COPY_BW=1.8TBs,d2d:3.6TBs
 
 [Fujitsu ARM]
 CPU vendor=Fujitsu ARM
diff --git a/src/uct/cuda/cuda_ipc/cuda_ipc_iface.c b/src/uct/cuda/cuda_ipc/cuda_ipc_iface.c
index b72dc01efa6..35c9fd56514 100644
--- a/src/uct/cuda/cuda_ipc/cuda_ipc_iface.c
+++ b/src/uct/cuda/cuda_ipc/cuda_ipc_iface.c
@@ -201,7 +201,7 @@ uct_cuda_ipc_iface_is_reachable_v2(const uct_iface_h tl_iface,
 static double uct_cuda_ipc_iface_get_bw()
 {
     CUdevice cu_device;
-    int major_version;
+    int major_version, minor_version = 0;
     ucs_status_t status;
 
     status = UCT_CUDADRV_FUNC_LOG_ERR(cuDeviceGet(&cu_device, 0));
@@ -217,6 +217,16 @@ static double uct_cuda_ipc_iface_get_bw()
         return 0;
     }
 
+    if (major_version == UCT_CUDA_BASE_GEN_B100) {
+        status = UCT_CUDADRV_FUNC_LOG_ERR(
+                cuDeviceGetAttribute(&minor_version,
+                                     CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MINOR,
+                                     cu_device));
+        if (status != UCS_OK) {
+            return 0;
+        }
+    }
+
     /*
      * TODO: Detect nvswitch
      */
@@ -230,6 +240,10 @@ static double uct_cuda_ipc_iface_get_bw()
     case UCT_CUDA_BASE_GEN_H100:
         return 400000.0 * UCS_MBYTE;
     case UCT_CUDA_BASE_GEN_B100:
+        if (minor_version == 7) {
+            return 1800000.0 * UCS_MBYTE;
+        }
+
         return 800000.0 * UCS_MBYTE;
     default:
         return 6911.0  * UCS_MBYTE;