TL/CUDA: Add int datatype support for NVLS

Juee14Desai · Juee14Desai · commit a6f32784bf28 · 2026-02-19T04:39:03.000+02:00
This PR adds support for the following additional integer data types in NVLS (NVLink SHARP) collective operations allreduce and reduce_scatter:

- INT32 (s32): 32-bit signed integer with v4 vectorization

- INT64 (s64): 64-bit signed integer with v2 vectorization

- UINT32 (u32): 32-bit unsigned integer with v4 vectorization

- UINT64 (u64): 64-bit unsigned integer with v2 vectorization

Added PTX multimem.ld_reduce and multimem.st instructions for each type

Created NvlsOps structs for type-specific operations

Updated allreduce and reduce_scatter kernels with new data type handling

Modified validation logic to accept new data types

Signed-off-by: Juee14Desai &lt;jueehimalbha@nvidia.com&gt;
diff --git a/src/components/tl/cuda/allreduce/allreduce_nvls.c b/src/components/tl/cuda/allreduce/allreduce_nvls.c
@@ -187,11 +187,15 @@ ucc_status_t ucc_tl_cuda_allreduce_nvls_init(
 
     if (buf_size < 1024 || coll_args->args.op != UCC_OP_SUM ||
         (coll_args->args.dst.info.datatype != UCC_DT_FLOAT32 &&
-         coll_args->args.dst.info.datatype != UCC_DT_BFLOAT16)) {
+         coll_args->args.dst.info.datatype != UCC_DT_BFLOAT16 &&
+         coll_args->args.dst.info.datatype != UCC_DT_INT32 &&
+         coll_args->args.dst.info.datatype != UCC_DT_UINT32 &&
+         coll_args->args.dst.info.datatype != UCC_DT_INT64 &&
+         coll_args->args.dst.info.datatype != UCC_DT_UINT64)) {
         tl_debug(
             UCC_TL_TEAM_LIB(team),
             "NVLS allreduce is supported only with SUM operation "
-            "and float32 or bfloat16 datatype, with message size >= 1024 "
+            "and float32, bfloat16, int32, uint32, int64, or uint64 datatype, with message size >= 1024 "
             "bytes");
         return UCC_ERR_NOT_SUPPORTED;
     }
diff --git a/src/components/tl/cuda/kernels/allreduce_kernel.cu b/src/components/tl/cuda/kernels/allreduce_kernel.cu
@@ -48,6 +48,72 @@ __global__ void __launch_bounds__(UCC_TL_CUDA_MAX_NVLS_THREADS)
     nvls_bar(&(mc_bar->arrival_counter), &(uc_bar->arrival_counter), total_blocks * (launch_counter * 2 + 2));
 }
 
+template <typename NvlsOps>
+__global__ void __launch_bounds__(UCC_TL_CUDA_MAX_NVLS_THREADS)
+    allreduce_kernel_scalar32(ucc_tl_cuda_nvls_control_t *mc_bar,
+                              ucc_tl_cuda_nvls_control_t *uc_bar,
+                              const uint32_t total_blocks,
+                              uint64_t launch_counter,
+                              uint32_t *base_u32, size_t count_u32, uint32_t rank,
+                              uint32_t tsize)
+{
+    // pre barrier
+    nvls_bar(&(mc_bar->arrival_counter), &(uc_bar->arrival_counter), total_blocks * (launch_counter * 2 + 1));
+
+    // Kernel execution
+    size_t chunk_start = ((int64_t)count_u32 * (int64_t)rank) / (int64_t)tsize;
+    size_t chunk_end   = ((int64_t)count_u32 * (int64_t)(rank + 1)) / (int64_t)tsize;
+
+    size_t thread_offset = (threadIdx.x + blockIdx.x * blockDim.x) * 4;
+    size_t stride        = blockDim.x * gridDim.x * 4;
+
+    for (size_t idx = chunk_start + thread_offset; idx < chunk_end; idx += stride) {
+        typename NvlsOps::value_type v0, v1, v2, v3;
+        NvlsOps::ld(v0, base_u32 + idx + 0);
+        NvlsOps::ld(v1, base_u32 + idx + 1);
+        NvlsOps::ld(v2, base_u32 + idx + 2);
+        NvlsOps::ld(v3, base_u32 + idx + 3);
+        NvlsOps::st(v0, base_u32 + idx + 0);
+        NvlsOps::st(v1, base_u32 + idx + 1);
+        NvlsOps::st(v2, base_u32 + idx + 2);
+        NvlsOps::st(v3, base_u32 + idx + 3);
+    }
+
+    // post barrier
+    nvls_bar(&(mc_bar->arrival_counter), &(uc_bar->arrival_counter), total_blocks * (launch_counter * 2 + 2));
+}
+
+template <typename NvlsOps>
+__global__ void __launch_bounds__(UCC_TL_CUDA_MAX_NVLS_THREADS)
+    allreduce_kernel_scalar64(ucc_tl_cuda_nvls_control_t *mc_bar,
+                              ucc_tl_cuda_nvls_control_t *uc_bar,
+                              const uint32_t total_blocks,
+                              uint64_t launch_counter,
+                              uint64_t *base_u64, size_t count_u64, uint32_t rank,
+                              uint32_t tsize)
+{
+    // pre barrier
+    nvls_bar(&(mc_bar->arrival_counter), &(uc_bar->arrival_counter), total_blocks * (launch_counter * 2 + 1));
+
+    // Kernel execution
+    size_t chunk_start = ((int64_t)count_u64 * (int64_t)rank) / (int64_t)tsize;
+    size_t chunk_end   = ((int64_t)count_u64 * (int64_t)(rank + 1)) / (int64_t)tsize;
+
+    size_t thread_offset = (threadIdx.x + blockIdx.x * blockDim.x) * 2;
+    size_t stride        = blockDim.x * gridDim.x * 2;
+
+    for (size_t idx = chunk_start + thread_offset; idx < chunk_end; idx += stride) {
+        typename NvlsOps::value_type v0, v1;
+        NvlsOps::ld(v0, base_u64 + idx + 0);
+        NvlsOps::ld(v1, base_u64 + idx + 1);
+        NvlsOps::st(v0, base_u64 + idx + 0);
+        NvlsOps::st(v1, base_u64 + idx + 1);
+    }
+
+    // post barrier
+    nvls_bar(&(mc_bar->arrival_counter), &(uc_bar->arrival_counter), total_blocks * (launch_counter * 2 + 2));
+}
+
 #ifdef __cplusplus
 extern "C" {
 #endif
@@ -69,17 +135,40 @@ ucc_status_t post_allreduce_kernel(cudaStream_t stream, uint32_t sm_count,
     ucc_tl_cuda_nvls_control_t *uc_bar = reinterpret_cast<ucc_tl_cuda_nvls_control_t *>(uc_control_addr);
     uint32_t expected_blocks = sm_count * tsize; // total num of blocks in the multicast group, num gpus * num blocks per gpu, used for barrier synchronization
 
+    assert(((uintptr_t)(mc_base_addr) % 8) == 0);
     switch (datatype) {
     case UCC_DT_FLOAT32:
-        assert(((uintptr_t)(mc_base_addr) % 8) == 0);
         allreduce_kernel_vec32<NvlsFp32Ops><<<sm_count, threads, 0, stream>>>(
             mc_bar, uc_bar, expected_blocks, launch_counter, base_u32, count_u32, rank, tsize);
         break;
     case UCC_DT_BFLOAT16:
-        assert(((uintptr_t)(mc_base_addr) % 8) == 0);
         allreduce_kernel_vec32<NvlsBf16Ops><<<sm_count, threads, 0, stream>>>(
             mc_bar, uc_bar, expected_blocks, launch_counter, base_u32, count_u32, rank, tsize);
         break;
+    case UCC_DT_INT32:
+        allreduce_kernel_scalar32<NvlsInt32Ops><<<sm_count, threads, 0, stream>>>(
+            mc_bar, uc_bar, expected_blocks, launch_counter, base_u32, count_u32, rank, tsize);
+        break;
+    case UCC_DT_UINT32:
+        allreduce_kernel_scalar32<NvlsUint32Ops><<<sm_count, threads, 0, stream>>>(
+            mc_bar, uc_bar, expected_blocks, launch_counter, base_u32, count_u32, rank, tsize);
+        break;
+    case UCC_DT_INT64:
+        {
+            uint64_t *base_u64 = reinterpret_cast<uint64_t *>(mc_base_addr);
+            size_t count_u64 = src_size_bytes / sizeof(uint64_t);
+            allreduce_kernel_scalar64<NvlsInt64Ops><<<sm_count, threads, 0, stream>>>(
+                mc_bar, uc_bar, expected_blocks, launch_counter, base_u64, count_u64, rank, tsize);
+        }
+        break;
+    case UCC_DT_UINT64:
+        {
+            uint64_t *base_u64 = reinterpret_cast<uint64_t *>(mc_base_addr);
+            size_t count_u64 = src_size_bytes / sizeof(uint64_t);
+            allreduce_kernel_scalar64<NvlsUint64Ops><<<sm_count, threads, 0, stream>>>(
+                mc_bar, uc_bar, expected_blocks, launch_counter, base_u64, count_u64, rank, tsize);
+        }
+        break;
     default:
         return UCC_ERR_NOT_SUPPORTED;
     }
diff --git a/src/components/tl/cuda/kernels/nvls.cuh b/src/components/tl/cuda/kernels/nvls.cuh
@@ -81,6 +81,67 @@ struct NvlsBf16Ops {
         MULTIMEM_ST_BF16(v, ptr);
     }
 };
+
+struct NvlsInt32Ops {
+    using value_type = unsigned int;
+    __device__ static inline void ld(unsigned int &v, const uint32_t *ptr) {
+        asm("multimem.ld_reduce.global.add.u32 %0, [%1];"
+            : "=r"(v)
+            : "l"(ptr)
+            : "memory");
+    }
+    __device__ static inline void st(const unsigned int &v, uint32_t *ptr) {
+        asm volatile("multimem.st.global.u32 [%0], %1;" ::"l"(ptr),
+                     "r"(v)
+                     : "memory");
+    }
+};
+
+struct NvlsUint32Ops {
+    using value_type = unsigned int;
+    __device__ static inline void ld(unsigned int &v, const uint32_t *ptr) {
+        asm("multimem.ld_reduce.global.add.u32 %0, [%1];"
+            : "=r"(v)
+            : "l"(ptr)
+            : "memory");
+    }
+    __device__ static inline void st(const unsigned int &v, uint32_t *ptr) {
+        asm volatile("multimem.st.global.u32 [%0], %1;" ::"l"(ptr),
+                     "r"(v)
+                     : "memory");
+    }
+};
+
+// PTX does not support s64 with add operation, so we use u64 instead
+struct NvlsInt64Ops {
+    using value_type = unsigned long long;
+    __device__ static inline void ld(unsigned long long &v, const uint64_t *ptr) {
+        asm("multimem.ld_reduce.global.add.u64 %0, [%1];"
+            : "=l"(v)
+            : "l"(ptr)
+            : "memory");
+    }
+    __device__ static inline void st(const unsigned long long &v, uint64_t *ptr) {
+        asm volatile("multimem.st.global.u64 [%0], %1;" ::"l"(ptr),
+                     "l"(v)
+                     : "memory");
+    }
+};
+
+struct NvlsUint64Ops {
+    using value_type = unsigned long long;
+    __device__ static inline void ld(unsigned long long &v, const uint64_t *ptr) {
+        asm("multimem.ld_reduce.global.add.u64 %0, [%1];"
+            : "=l"(v)
+            : "l"(ptr)
+            : "memory");
+    }
+    __device__ static inline void st(const unsigned long long &v, uint64_t *ptr) {
+        asm volatile("multimem.st.global.u64 [%0], %1;" ::"l"(ptr),
+                     "l"(v)
+                     : "memory");
+    }
+};
 #endif // __cplusplus
 
 #endif // UCC_TL_CUDA_NVLS_CUH_
diff --git a/src/components/tl/cuda/kernels/reduce_scatter_kernel.cu b/src/components/tl/cuda/kernels/reduce_scatter_kernel.cu
@@ -47,6 +47,74 @@ __global__ void __launch_bounds__(UCC_TL_CUDA_MAX_NVLS_THREADS)
         total_blocks * (launch_counter * 2 + 2));
 }
 
+template <typename NvlsOps>
+__global__ void __launch_bounds__(UCC_TL_CUDA_MAX_NVLS_THREADS)
+    reduce_scatter_kernel_scalar32(
+        ucc_tl_cuda_nvls_control_t *mc_bar, ucc_tl_cuda_nvls_control_t *uc_bar,
+        const uint32_t total_blocks, uint64_t launch_counter,
+        uint32_t *base_u32, size_t offset, size_t count, uint32_t *dst_u32)
+{
+    // pre barrier
+    nvls_bar(
+        &(mc_bar->arrival_counter),
+        &(uc_bar->arrival_counter),
+        total_blocks * (launch_counter * 2 + 1));
+
+    size_t thread_offset = (threadIdx.x + blockIdx.x * blockDim.x) * 4;
+    size_t stride        = blockDim.x * gridDim.x * 4;
+
+    for (size_t idx = offset + thread_offset; idx < offset + count;
+         idx += stride) {
+        typename NvlsOps::value_type v0, v1, v2, v3;
+        NvlsOps::ld(v0, base_u32 + idx + 0);
+        NvlsOps::ld(v1, base_u32 + idx + 1);
+        NvlsOps::ld(v2, base_u32 + idx + 2);
+        NvlsOps::ld(v3, base_u32 + idx + 3);
+        dst_u32[idx - offset + 0] = v0;
+        dst_u32[idx - offset + 1] = v1;
+        dst_u32[idx - offset + 2] = v2;
+        dst_u32[idx - offset + 3] = v3;
+    }
+
+    // post barrier
+    nvls_bar(
+        &(mc_bar->arrival_counter),
+        &(uc_bar->arrival_counter),
+        total_blocks * (launch_counter * 2 + 2));
+}
+
+template <typename NvlsOps>
+__global__ void __launch_bounds__(UCC_TL_CUDA_MAX_NVLS_THREADS)
+    reduce_scatter_kernel_scalar64(
+        ucc_tl_cuda_nvls_control_t *mc_bar, ucc_tl_cuda_nvls_control_t *uc_bar,
+        const uint32_t total_blocks, uint64_t launch_counter,
+        uint64_t *base_u64, size_t offset, size_t count, uint64_t *dst_u64)
+{
+    // pre barrier
+    nvls_bar(
+        &(mc_bar->arrival_counter),
+        &(uc_bar->arrival_counter),
+        total_blocks * (launch_counter * 2 + 1));
+
+    size_t thread_offset = (threadIdx.x + blockIdx.x * blockDim.x) * 2;
+    size_t stride        = blockDim.x * gridDim.x * 2;
+
+    for (size_t idx = offset + thread_offset; idx < offset + count;
+         idx += stride) {
+        typename NvlsOps::value_type v0, v1;
+        NvlsOps::ld(v0, base_u64 + idx + 0);
+        NvlsOps::ld(v1, base_u64 + idx + 1);
+        dst_u64[idx - offset + 0] = v0;
+        dst_u64[idx - offset + 1] = v1;
+    }
+
+    // post barrier
+    nvls_bar(
+        &(mc_bar->arrival_counter),
+        &(uc_bar->arrival_counter),
+        total_blocks * (launch_counter * 2 + 2));
+}
+
 #ifdef __cplusplus
 extern "C" {
 #endif
@@ -101,6 +169,64 @@ ucc_status_t post_reduce_scatter_kernel(
                 count,
                 reinterpret_cast<uint32_t *>(dst_ptr));
         break;
+    case UCC_DT_INT32:
+        reduce_scatter_kernel_scalar32<NvlsInt32Ops>
+            <<<sm_count, threads, 0, stream>>>(
+                mc_bar,
+                uc_bar,
+                expected_blocks,
+                launch_counter,
+                base_u32,
+                offset,
+                count,
+                reinterpret_cast<uint32_t *>(dst_ptr));
+        break;
+    case UCC_DT_UINT32:
+        reduce_scatter_kernel_scalar32<NvlsUint32Ops>
+            <<<sm_count, threads, 0, stream>>>(
+                mc_bar,
+                uc_bar,
+                expected_blocks,
+                launch_counter,
+                base_u32,
+                offset,
+                count,
+                reinterpret_cast<uint32_t *>(dst_ptr));
+        break;
+    case UCC_DT_INT64:
+        {
+            uint64_t *base_u64 = reinterpret_cast<uint64_t *>(mc_base_addr);
+            size_t offset_u64 = offset / 2;
+            size_t count_u64 = count / 2;
+            reduce_scatter_kernel_scalar64<NvlsInt64Ops>
+                <<<sm_count, threads, 0, stream>>>(
+                    mc_bar,
+                    uc_bar,
+                    expected_blocks,
+                    launch_counter,
+                    base_u64,
+                    offset_u64,
+                    count_u64,
+                    reinterpret_cast<uint64_t *>(dst_ptr));
+        }
+        break;
+    case UCC_DT_UINT64:
+        {
+            uint64_t *base_u64 = reinterpret_cast<uint64_t *>(mc_base_addr);
+            size_t offset_u64 = offset / 2;
+            size_t count_u64 = count / 2;
+            reduce_scatter_kernel_scalar64<NvlsUint64Ops>
+                <<<sm_count, threads, 0, stream>>>(
+                    mc_bar,
+                    uc_bar,
+                    expected_blocks,
+                    launch_counter,
+                    base_u64,
+                    offset_u64,
+                    count_u64,
+                    reinterpret_cast<uint64_t *>(dst_ptr));
+        }
+        break;
     default:
         return UCC_ERR_NOT_SUPPORTED;
     }
diff --git a/src/components/tl/cuda/reduce_scatterv/reduce_scatterv_nvls.c b/src/components/tl/cuda/reduce_scatterv/reduce_scatterv_nvls.c
@@ -175,10 +175,12 @@ ucc_status_t ucc_tl_cuda_reduce_scatterv_nvls_init_common(
             "NVLS reduce scatter(v) supported only with SUM operation");
         return UCC_ERR_NOT_SUPPORTED;
     }
-    if (dt != UCC_DT_FLOAT32 && dt != UCC_DT_BFLOAT16) {
+    if (dt != UCC_DT_FLOAT32 && dt != UCC_DT_BFLOAT16 &&
+        dt != UCC_DT_INT32 && dt != UCC_DT_UINT32 &&
+        dt != UCC_DT_INT64 && dt != UCC_DT_UINT64) {
         tl_debug(
             UCC_TL_TEAM_LIB(team),
-            "NVLS reduce scatter(v) supported only with float32/bfloat16");
+            "NVLS reduce scatter(v) supported only with float32/bfloat16/int32/uint32/int64/uint64");
         return UCC_ERR_NOT_SUPPORTED;
     }
 
@@ -191,11 +193,15 @@ ucc_status_t ucc_tl_cuda_reduce_scatterv_nvls_init_common(
     count_elements  = get_count(task, trank);
 
     /* Convert from datatype elements to uint32_t indices for the kernel.
-     * For float32: 1 element = 1 uint32_t
+     * For float32/int32/uint32: 1 element = 1 uint32_t
+     * For int64/uint64: 1 element = 2 uint32_t
      * For bfloat16: 2 elements = 1 uint32_t */
-    if (dt == UCC_DT_FLOAT32) {
+    if (dt == UCC_DT_FLOAT32 || dt == UCC_DT_INT32 || dt == UCC_DT_UINT32) {
         offset_u32 = offset_elements;
         count_u32  = count_elements;
+    } else if (dt == UCC_DT_INT64 || dt == UCC_DT_UINT64) {
+        offset_u32 = offset_elements * 2;
+        count_u32  = count_elements * 2;
     } else { /* UCC_DT_BFLOAT16 */
         if (offset_elements % 2 != 0 || count_elements % 2 != 0) {
             tl_debug(
