issue/434 nccl support bf16

81093e0b · PanZezhong1725 · 9ad23fad · 81093e0b
Commit 81093e0b authored Sep 09, 2025 by PanZezhong1725
Show whitespace changes
Inline Side-by-side

Showing with 3 additions and 3 deletions

src/infiniccl/cuda/infiniccl_cuda.cu src/infiniccl/cuda/infiniccl_cuda.cu +3 -3

No files found.
--- a/src/infiniccl/cuda/infiniccl_cuda.cu
+++ b/src/infiniccl/cuda/infiniccl_cuda.cu
@@ -22,6 +22,8 @@ inline ncclDataType_t getNcclDtype(infiniDtype_t datatype) {
        return ncclFloat;
    case INFINI_DTYPE_F16:
        return ncclHalf;
+    case INFINI_DTYPE_BF16:
+        return ncclBfloat16;
    default:
        std::abort();
        return ncclHalf;
@@ -82,9 +84,7 @@ infiniStatus_t allReduce(
    infinicclComm_t comm,
    infinirtStream_t stream) {
-    if (datatype != INFINI_DTYPE_F32 && datatype != INFINI_DTYPE_F16) {
+    CHECK_DTYPE(datatype, INFINI_DTYPE_F32, INFINI_DTYPE_F16, INFINI_DTYPE_BF16);
-        return INFINI_STATUS_BAD_PARAM;
-    }
    CHECK_NCCL(ncclAllReduce(sendbuf, recvbuf, count, getNcclDtype(datatype),
                             getNcclRedOp(op), getNcclComm(comm), getCudaStream(stream)));