[DeepEP] support M2N

paddle/fluid/distributed/collective/deep_ep/CMakeLists.txt

@@ -7,8 +7,13 @@ if(WITH_NVSHMEM)
                        CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS}")
 
   set(DEEPEP_KERNEL_SRCS
-      kernels/intranode.cu kernels/runtime.cu kernels/internode.cu
-      kernels/internode_ll.cu kernels/internode_ll_two_stage.cu)
+      kernels/intranode.cu
+      kernels/runtime.cu
+      kernels/internode.cu
+      kernels/internode_ll.cu
+      kernels/internode_ll_two_stage.cu
+      kernels/internode_ll.cu
+      kernels/m2n_ll_two_stage.cu)
   cc_library(
     deepep_kernels
     SRCS ${DEEPEP_KERNEL_SRCS}

paddle/fluid/distributed/collective/deep_ep/config.hpp

@@ -149,10 +149,14 @@ struct LowLatencyBuffer {
   void* dispatch_rdma_send_buffer = nullptr;
   void* dispatch_rdma_recv_data_buffer = nullptr;
   int* dispatch_rdma_recv_count_buffer = nullptr;
+  // Note(ZKK) this is only used in M2N !
+  int* dispatch_rdma_recv_complete_buffer = nullptr;
 
   void* combine_rdma_send_buffer = nullptr;
   void* combine_rdma_recv_data_buffer = nullptr;
   int* combine_rdma_recv_flag_buffer = nullptr;
+  // Note(ZKK) this is only used in M2N !
+  int* combine_rdma_recv_complete_buffer = nullptr;
 
   void* combine_rdma_send_buffer_data_start = nullptr;
   size_t num_bytes_per_combine_msg = 0;
@@ -244,11 +248,19 @@ struct LowLatencyLayout {
           advance<int*>(rdma_buffer,
                         send_buffer_bytes * 2 + recv_buffer_bytes * 2 +
                             signaling_buffer_bytes * i),
+          // Note(ZKK): dispatch_rdma_recv_complete_buffer is only used in M2N!
+          // so here we symbolically add a 0 to it
+          advance<int*>(rdma_buffer, 0),
+
           advance(rdma_buffer, send_buffer_bytes * i),
           advance(rdma_buffer, send_buffer_bytes * 2 + recv_buffer_bytes * i),
           advance<int*>(rdma_buffer,
                         send_buffer_bytes * 2 + recv_buffer_bytes * 2 +
                             signaling_buffer_bytes * i),
+          // Note(ZKK): combine_rdma_recv_complete_buffer is only used in M2N!
+          // so here we symbolically add a 0 to it
+          advance<int*>(rdma_buffer, 0),
+
           advance(rdma_buffer, send_buffer_bytes * i),
           num_bytes_per_combine_msg};
     }
@@ -318,6 +330,12 @@ struct LowLatencyTwoStageLayout {
                                              combine_recv_flag_buffer_bytes);
     total_bytes += signaling_buffer_bytes * 2;
 
+    // Symmetric complete signaling buffers
+    // Note(ZKK): this is only used in M2N!
+    size_t recv_complete_buffer_bytes =
+        2 * M2N_NUM_MAX_MICRO_BATCHES * num_ranks * sizeof(int);
+    total_bytes += recv_complete_buffer_bytes * 2;
+
     // Assign pointers
     for (int i = 0; i < 2; ++i) {
       buffers[i] = {
@@ -327,11 +345,21 @@ struct LowLatencyTwoStageLayout {
           advance<int*>(rdma_buffer,
                         send_buffer_bytes * 2 + recv_buffer_bytes * 2 +
                             signaling_buffer_bytes * i),
+          // dispatch_rdma_recv_complete_buffer!
+          advance<int*>(rdma_buffer,
+                        send_buffer_bytes * 2 + recv_buffer_bytes * 2 +
+                            signaling_buffer_bytes * 2 +
+                            recv_complete_buffer_bytes * i),
           advance(rdma_buffer, send_buffer_bytes * i),
           advance(rdma_buffer, send_buffer_bytes * 2 + recv_buffer_bytes * i),
           advance<int*>(rdma_buffer,
                         send_buffer_bytes * 2 + recv_buffer_bytes * 2 +
                             signaling_buffer_bytes * i),
+          // combine_rdma_recv_complete_buffer!
+          advance<int*>(rdma_buffer,
+                        send_buffer_bytes * 2 + recv_buffer_bytes * 2 +
+                            signaling_buffer_bytes * 2 +
+                            recv_complete_buffer_bytes * i),
           advance(rdma_buffer, send_buffer_bytes * i),
           num_bytes_per_combine_msg};
     }