[PAL] Implement proper build-time detection and filtering of NVCC gencodes to enable NVIDIA multicast primitive support (#387)

Author: MC952-arch

Makefile

@@ -144,6 +144,9 @@ ifeq ($(USE_NVIDIA), 1)
 	CCL_INCLUDE = $(CCL_HOME)/include
 	CCL_LINK = -lnccl
 	ADAPTOR_FLAG = -DUSE_NVIDIA_ADAPTOR
+ifeq ($(NVCC_GENCODE_MULTICAST_UNSUPPORTED), 1)
+	ADAPTOR_FLAG += -DNVCC_GENCODE_MULTICAST_UNSUPPORTED
+endif
 else ifeq ($(USE_ASCEND), 1)
 	DEVICE_LIB = $(DEVICE_HOME)/lib64
 	DEVICE_INCLUDE = $(DEVICE_HOME)/include

flagcx/adaptor/ccl/nccl_adaptor.cc

@@ -332,7 +332,8 @@ flagcxResult_t ncclAdaptorAllReduce(const void *sendbuff, void *recvbuff,
                                     flagcxRedOp_t op, flagcxInnerComm_t comm,
                                     flagcxStream_t stream) {
 #if defined(COMPILE_KERNEL_HOST) &&                                            \
-    (NCCL_VERSION_CODE > NCCL_VERSION(2, 28, 0)) && (__CUDA_ARCH__ >= 900)
+    (NCCL_VERSION_CODE > NCCL_VERSION(2, 28, 0)) &&                            \
+    !defined(NVCC_GENCODE_MULTICAST_UNSUPPORTED)
   size_t size = count * getFlagcxDataTypeSize(datatype);
   int nranks;
   FLAGCXCHECK((flagcxResult_t)ncclCommCount(comm->base, &nranks));
@@ -362,8 +363,7 @@ flagcxResult_t ncclAdaptorAllReduce(const void *sendbuff, void *recvbuff,
   FLAGCXCHECK((flagcxResult_t)ncclAllReduce(
       sendbuff, recvbuff, count, (ncclDataType_t)datatype, (ncclRedOp_t)op,
       comm->base, stream->base));
-#endif // defined(COMPILE_KERNEL_HOST) && (NCCL_VERSION_CODE > NCCL_VERSION(2,
-       // 28, 0) && (__CUDA_ARCH__ >= 900))
+#endif
   return flagcxSuccess;
 }
 

flagcx/kernels/custom_allreduce.cu

@@ -148,6 +148,7 @@ FLAGCX_DEVICE_INLINE_DECORATOR typename packed_t<T, ByteSize>::array_type
 multimem_sum(T* addr) {
   using P = packed_t<T, ByteSize>;
   typename P::array_type ret;
+#if __CUDA_ARCH__ >= 900
   if constexpr (std::is_same<T, nv_bfloat16>::value) {
     typename P::storage_t h;
     asm volatile(
@@ -177,6 +178,7 @@ multimem_sum(T* addr) {
         : "l"(addr)
         : "memory");
   }
+#endif
   return ret;
 }
 
@@ -188,6 +190,7 @@ multimem_min(T* addr) {
                 "multimem min only supports half and bfloat16");
   using P = packed_t<T, ByteSize>;
   typename P::array_type ret;
+#if __CUDA_ARCH__ >= 900
   if constexpr (std::is_same<T, nv_bfloat16>::value) {
     typename P::storage_t h;
     asm volatile(
@@ -205,6 +208,7 @@ multimem_min(T* addr) {
         : "memory");
     unpack<T, ByteSize>(h, ret.data);
   }
+#endif
   return ret;
 }
 
@@ -216,6 +220,7 @@ multimem_max(T* addr) {
                 "multimem max only supports half and bfloat16");
   using P = packed_t<T, ByteSize>;
   typename P::array_type ret;
+#if __CUDA_ARCH__ >= 900
   if constexpr (std::is_same<T, nv_bfloat16>::value) {
     typename P::storage_t h;
     asm volatile(
@@ -233,6 +238,7 @@ multimem_max(T* addr) {
         : "memory");
     unpack<T, ByteSize>(h, ret.data);
   }
+#endif
   return ret;
 }
 
@@ -264,6 +270,7 @@ multimem_reduce(T* addr, ncclRedOp_t op) {
 template <typename T, int ByteSize = (sizeof(T) <= 4 ? 4 : 8)>
 FLAGCX_DEVICE_INLINE_DECORATOR void
 multimem_st(T* addr, typename packed_t<T, ByteSize>::array_type val) {
+#if __CUDA_ARCH__ >= 900
   using P = packed_t<T, ByteSize>;
   if constexpr (std::is_same<T, nv_bfloat16>::value) {
     typename P::storage_t h = pack<T, ByteSize>(val.data);
@@ -292,6 +299,7 @@ multimem_st(T* addr, typename packed_t<T, ByteSize>::array_type val) {
         : "l"(addr), "d"(val.data[0])
         : "memory");
   }
+#endif
 }
 
 // Store to local/shared memory using vectorized store

makefiles/nvidia_gencode.mk

@@ -88,4 +88,32 @@ else
     $(CUDA8_GENCODE) \
     $(CUDA8_PTX)
 endif
-# $(info DEVICE_COMPILER_GENCODE is ${DEVICE_COMPILER_GENCODE})
+
+# ----------------
+# Multicast primitive support check (requires sm_90+)
+# If sm_90+ gencode exists: filter out < sm_90 gencodes
+# If no sm_90+ gencode: keep all gencodes and define NVCC_GENCODE_MULTICAST_UNSUPPORTED
+# ----------------
+
+# Define all sm_90+ gencodes for detection
+SM90_PLUS_GENCODES = $(CUDA12_GENCODE) $(CUDA12_8_GENCODE) $(CUDA13_GENCODE) $(CUDA12_PTX) $(CUDA13_PTX)
+
+# Check if any sm_90+ gencode is present in DEVICE_COMPILER_GENCODE
+HAS_SM90_PLUS := $(strip $(foreach gc,$(SM90_PLUS_GENCODES),$(findstring $(gc),$(DEVICE_COMPILER_GENCODE))))
+
+ifneq ($(HAS_SM90_PLUS),)
+  # Has sm_90+ support: filter out older gencodes for multicast kernel compilation
+  DEVICE_COMPILER_GENCODE := $(filter-out $(CUDA8_GENCODE),$(DEVICE_COMPILER_GENCODE))
+  DEVICE_COMPILER_GENCODE := $(filter-out $(CUDA9_GENCODE),$(DEVICE_COMPILER_GENCODE))
+  DEVICE_COMPILER_GENCODE := $(filter-out $(CUDA10_GENCODE),$(DEVICE_COMPILER_GENCODE))
+  DEVICE_COMPILER_GENCODE := $(filter-out $(CUDA11_GENCODE),$(DEVICE_COMPILER_GENCODE))
+  DEVICE_COMPILER_GENCODE := $(filter-out $(CUDA8_PTX),$(DEVICE_COMPILER_GENCODE))
+  DEVICE_COMPILER_GENCODE := $(filter-out $(CUDA9_PTX),$(DEVICE_COMPILER_GENCODE))
+  DEVICE_COMPILER_GENCODE := $(filter-out $(CUDA11_PTX),$(DEVICE_COMPILER_GENCODE))
+  $(info Multicast support enabled. Using DEVICE_COMPILER_GENCODE: ${DEVICE_COMPILER_GENCODE})
+else
+  # No sm_90+ support: keep all gencodes and disable multicast
+  NVCC_GENCODE_MULTICAST_UNSUPPORTED := 1
+  $(info WARNING: No sm_90+ architecture detected. Multicast primitives disabled. Custom AllReduce will fallback to ncclAllReduce.)
+  $(info Using DEVICE_COMPILER_GENCODE: ${DEVICE_COMPILER_GENCODE})
+endif