refine

Oneflow-Inc · Jan 3, 2025 · ad9e7ee · ad9e7ee
1 parent e5777b9
commit ad9e7ee
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 9 deletions.
diff --git a/oneflow/core/kernel/nccl_send_recv_boxing_kernel.cpp b/oneflow/core/kernel/nccl_send_recv_boxing_kernel.cpp
@@ -23,6 +23,8 @@ limitations under the License.
 #include "oneflow/user/kernels/collective_communication/include/send.h"
 #include "oneflow/user/kernels/collective_communication/include/recv.h"
 
+#if defined(WITH_CUDA) && NCCL_VERSION_CODE > 2700
+
 namespace oneflow {
 
 class NcclSendRecvBoxingKernel final : public Kernel {
@@ -255,3 +257,5 @@ REGISTER_KERNEL(OperatorConf::kNcclSendRecvBoxingConf, NcclSendRecvBoxingKernel)
 REGISTER_SYSTEM_OP_KERNEL_UNIFIED_NCCL_COMM_INIT(OperatorConf::kNcclSendRecvBoxingConf);
 
 }  // namespace oneflow
+
+#endif  // WITH_CUDA && NCCL_VERSION_CODE > 2700
diff --git a/oneflow/user/kernels/collective_communication/include/collective_communication.h b/oneflow/user/kernels/collective_communication/include/collective_communication.h
@@ -16,8 +16,6 @@ limitations under the License.
 #ifndef ONEFLOW_USER_KERNELS_COLLECTIVE_COMMUNICATION_INCLUDE_COLLECTIVE_COMMUNICATION_H_
 #define ONEFLOW_USER_KERNELS_COLLECTIVE_COMMUNICATION_INCLUDE_COLLECTIVE_COMMUNICATION_H_
 
-#include <memory>
-#include <utility>
 #include "oneflow/core/common/auto_registration_factory.h"
 #include "oneflow/core/common/switch_func.h"
 #include "oneflow/user/kernels/collective_communication/include/communication_context.h"

diff --git a/oneflow/user/kernels/nccl_logical_fusion_kernel.cpp b/oneflow/user/kernels/nccl_logical_fusion_kernel.cpp
@@ -13,20 +13,19 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
 */
-
-#include "collective_communication/include/all_gather.h"
-#include "collective_communication/include/all_reduce.h"
-#include "collective_communication/include/collective_communication.h"
-#include "collective_communication/include/reduce_scatter.h"
 #include "oneflow/core/framework/framework.h"
 #include "oneflow/core/device/nccl_util.h"
 #include "oneflow/core/job/eager_nccl_comm_manager.h"
 #include "oneflow/core/job/parallel_desc.h"
 #include "oneflow/core/ep/include/primitive/permute.h"
 #include "oneflow/core/ep/cuda/cuda_stream.h"
 #include "oneflow/user/ops/nccl_logical_util.h"
-#include "oneflow/user/kernels/collective_communication/include/send.h"
-#include "oneflow/user/kernels/collective_communication/include/recv.h"
+#include "collective_communication/include/send.h"
+#include "collective_communication/include/recv.h"
+#include "collective_communication/include/all_gather.h"
+#include "collective_communication/include/all_reduce.h"
+#include "collective_communication/include/collective_communication.h"
+#include "collective_communication/include/reduce_scatter.h"
 
 #if defined(WITH_CUDA) && NCCL_VERSION_CODE > 2700