Revert "Update deep_ep intranode & internode kernels (#74284)" (#76090)

Author: lshpku

paddle/fluid/distributed/collective/deep_ep/deep_ep.cpp

@@ -83,11 +83,10 @@ Buffer::Buffer(int rank,
   calc_ctx = reinterpret_cast<phi::GPUContext*>(
       reinterpret_cast<paddle::distributed::ProcessGroupNCCL*>(pg)
           ->GetDeviceContext(place, true));
-
-  // Metadata memory
-  int64_t barrier_signal_bytes = NUM_MAX_NVL_PEERS * sizeof(int);
-  int64_t buffer_ptr_bytes = NUM_MAX_NVL_PEERS * sizeof(void*);
-  int64_t barrier_signal_ptr_bytes = NUM_MAX_NVL_PEERS * sizeof(int*);
+  // Task fifo memory
+  int64_t fifo_bytes = sizeof(int) * NUM_MAX_FIFO_SLOTS;
+  int64_t buffer_ptr_bytes = sizeof(void*) * NUM_MAX_NVL_PEERS;
+  int64_t task_ptr_bytes = sizeof(int*) * NUM_MAX_NVL_PEERS;
 
   // Common checks
   EP_HOST_ASSERT(
@@ -106,35 +105,40 @@ Buffer::Buffer(int rank,
     EP_HOST_ASSERT(num_ranks > NUM_MAX_NVL_PEERS || low_latency_mode);
 
   // Get ranks
+  // CUDA_CHECK(cudaGetDevice(&device_id));
   rdma_rank = rank / NUM_MAX_NVL_PEERS, nvl_rank = rank % NUM_MAX_NVL_PEERS;
-  num_rdma_ranks = std::max(1, num_ranks / NUM_MAX_NVL_PEERS);
+  num_rdma_ranks = std::max(1, num_ranks / NUM_MAX_NVL_PEERS),
   num_nvl_ranks = std::min(num_ranks, NUM_MAX_NVL_PEERS);
 
   // Get device info
   cudaDeviceProp device_prop = {};
   CUDA_CHECK(cudaGetDeviceProperties(&device_prop, device_id));
 
   if (num_nvl_bytes > 0) {
-    // Local IPC: alloc local memory and set local IPC handles
-    CUDA_CHECK(cudaMalloc(&buffer_ptrs[nvl_rank],
-                          num_nvl_bytes + barrier_signal_bytes +
-                              buffer_ptr_bytes + barrier_signal_ptr_bytes));
+    // Local IPC: alloc local memory and set local IPC handle
+    CUDA_CHECK(cudaMalloc(
+        &buffer_ptrs[nvl_rank],
+        num_nvl_bytes + fifo_bytes + buffer_ptr_bytes + task_ptr_bytes));
     CUDA_CHECK(
         cudaIpcGetMemHandle(&ipc_handles[nvl_rank], buffer_ptrs[nvl_rank]));
-    buffer_ptrs_gpu =
-        reinterpret_cast<void**>(static_cast<uint8_t*>(buffer_ptrs[nvl_rank]) +
-                                 num_nvl_bytes + barrier_signal_bytes);
-
-    // Set barrier signals
-    barrier_signal_ptrs[nvl_rank] = reinterpret_cast<int*>(
-        static_cast<uint8_t*>(buffer_ptrs[nvl_rank]) + num_nvl_bytes);
-    barrier_signal_ptrs_gpu = reinterpret_cast<int**>(
-        static_cast<uint8_t*>(buffer_ptrs[nvl_rank]) + num_nvl_bytes +
-        barrier_signal_bytes + buffer_ptr_bytes);
+    buffer_ptrs_gpu = reinterpret_cast<void**>(
+        reinterpret_cast<uint8_t*>(buffer_ptrs[nvl_rank]) + num_nvl_bytes +
+        fifo_bytes);
+
+    // Set task fifo
+    EP_HOST_ASSERT(NUM_MAX_FIFO_SLOTS % num_nvl_ranks == 0);
+    task_fifo_ptrs[nvl_rank] = reinterpret_cast<int*>(
+        reinterpret_cast<uint8_t*>(buffer_ptrs[nvl_rank]) + num_nvl_bytes);
+    task_fifo_ptrs_gpu = reinterpret_cast<int**>(
+        reinterpret_cast<uint8_t*>(buffer_ptrs[nvl_rank]) + num_nvl_bytes +
+        fifo_bytes + buffer_ptr_bytes);
 
     // No need to synchronize, will do a full device sync during `sync`
     CUDA_CHECK(cudaMemsetAsync(
-        barrier_signal_ptrs[nvl_rank], 0, barrier_signal_bytes, comm_stream));
+        buffer_ptrs[nvl_rank],
+        0,
+        num_nvl_bytes + fifo_bytes + buffer_ptr_bytes + task_ptr_bytes,
+        comm_stream));
   }
 
   // Create 32 MiB workspace
@@ -180,7 +184,8 @@ Buffer::~Buffer() noexcept(false) {
   if (num_nvl_bytes > 0) {
     // Barrier
     intranode::barrier(
-        barrier_signal_ptrs_gpu, nvl_rank, num_nvl_ranks, comm_stream);
+        task_fifo_ptrs_gpu, head, nvl_rank, num_nvl_ranks, comm_stream);
+    move_fifo_slots();
     CUDA_CHECK(cudaDeviceSynchronize());
 
     // Close remote IPC
@@ -211,6 +216,10 @@ Buffer::~Buffer() noexcept(false) {
   CUDA_CHECK(cudaFreeHost(const_cast<int*>(moe_recv_expert_counter)));
 }
 
+void Buffer::move_fifo_slots(int num_slots) {
+  head = (head + num_ranks * num_slots) % NUM_MAX_FIFO_SLOTS;
+}
+
 bool Buffer::is_available() const { return available; }
 
 bool Buffer::is_internode_available() const {
@@ -259,7 +268,7 @@ void Buffer::sync(
 
   // Sync IPC handles
   if (num_nvl_bytes > 0) {
-    EP_HOST_ASSERT(num_ranks == device_ids.size());
+    EP_HOST_ASSERT(num_ranks == static_cast<int64_t>(device_ids.size()));
     EP_HOST_ASSERT(device_ids.size() == all_gathered_handles.size());
     for (int i = 0, offset = rdma_rank * num_nvl_ranks; i < num_nvl_ranks;
          ++i) {
@@ -271,22 +280,22 @@ void Buffer::sync(
             ipc_handles[i].reserved, handle_str.c_str(), CUDA_IPC_HANDLE_SIZE);
         CUDA_CHECK(cudaIpcOpenMemHandle(
             &buffer_ptrs[i], ipc_handles[i], cudaIpcMemLazyEnablePeerAccess));
-        barrier_signal_ptrs[i] = reinterpret_cast<int*>(
-            static_cast<uint8_t*>(buffer_ptrs[i]) + num_nvl_bytes);
+        task_fifo_ptrs[i] = reinterpret_cast<int*>(
+            reinterpret_cast<uint8_t*>(buffer_ptrs[i]) + num_nvl_bytes);
       } else {
         EP_HOST_ASSERT(std::memcmp(ipc_handles[i].reserved,
                                    handle_str.c_str(),
                                    CUDA_IPC_HANDLE_SIZE) == 0);
       }
     }
 
-    // Copy all buffer and barrier signal pointers to GPU
+    // Copy all buffer and task pointers to GPU
     CUDA_CHECK(cudaMemcpy(buffer_ptrs_gpu,
                           buffer_ptrs,
                           sizeof(void*) * NUM_MAX_NVL_PEERS,
                           cudaMemcpyHostToDevice));
-    CUDA_CHECK(cudaMemcpy(barrier_signal_ptrs_gpu,
-                          barrier_signal_ptrs,
+    CUDA_CHECK(cudaMemcpy(task_fifo_ptrs_gpu,
+                          task_fifo_ptrs,
                           sizeof(int*) * NUM_MAX_NVL_PEERS,
                           cudaMemcpyHostToDevice));
     CUDA_CHECK(cudaDeviceSynchronize());
@@ -530,7 +539,7 @@ Buffer::intranode_dispatch(
 
   // FP8 scales checks
   float* x_scales_ptr = nullptr;
-  int num_scales = 0, scale_token_stride = 0, scale_hidden_stride = 0;
+  int num_scales = 0;
   if (x_scales.has_value()) {
     EP_HOST_ASSERT(x.element_size() == 1);
     EP_HOST_ASSERT(x_scales->scalar_type() == deep_ep::detail::kFloat32);
@@ -539,8 +548,6 @@ Buffer::intranode_dispatch(
     EP_HOST_ASSERT(x_scales->size(0) == num_tokens);
     num_scales = x_scales->dim() == 1 ? 1 : static_cast<int>(x_scales->size(1));
     x_scales_ptr = x_scales->data_ptr<float>();
-    scale_token_stride = static_cast<int>(x_scales->stride(0));
-    scale_hidden_stride = static_cast<int>(x_scales->stride(1));
   }
 
   // Allocate all tensors on comm stream if set
@@ -579,10 +586,12 @@ Buffer::intranode_dispatch(
     intranode::cached_notify_dispatch(rank_prefix_matrix.data_ptr<int>(),
                                       num_memset_int,
                                       buffer_ptrs_gpu,
-                                      barrier_signal_ptrs_gpu,
+                                      task_fifo_ptrs_gpu,
+                                      head,
                                       rank,
                                       num_ranks,
                                       comm_stream);
+    move_fifo_slots(2);
   } else {
     rank_prefix_matrix = ConvertPaddleTensorToDetailTensor(
         paddle::experimental::empty({num_ranks, num_ranks},
@@ -617,10 +626,12 @@ Buffer::intranode_dispatch(
                                num_memset_int,
                                expert_alignment,
                                buffer_ptrs_gpu,
-                               barrier_signal_ptrs_gpu,
+                               task_fifo_ptrs_gpu,
+                               head,
                                rank,
                                comm_stream,
                                num_channels);
+    move_fifo_slots(3);
 
     // Synchronize total received tokens and tokens per expert
     auto start_time = std::chrono::high_resolution_clock::now();
@@ -730,13 +741,10 @@ Buffer::intranode_dispatch(
       is_token_in_rank.data_ptr<bool>(),
       channel_prefix_matrix.data_ptr<int>(),
       num_tokens,
-      0,  // num_worst_tokens (not exposed)
       static_cast<int>(hidden * recv_x.element_size() / sizeof(int4)),
       num_topk,
       num_experts,
       num_scales,
-      scale_token_stride,
-      scale_hidden_stride,
       buffer_ptrs_gpu,
       rank,
       num_ranks,
@@ -881,11 +889,15 @@ Buffer::intranode_combine(
                                    num_channels,
                                    num_recv_tokens,
                                    num_channels * num_ranks * 2,
-                                   barrier_signal_ptrs_gpu,
+                                   task_fifo_ptrs_gpu,
+                                   head,
                                    rank,
                                    num_ranks,
                                    comm_stream);
 
+  // NOTES: this function uses two FIFO slots (barrier before and after)
+  move_fifo_slots(2);
+
   // Combine data
   auto recv_x = ConvertPaddleTensorToDetailTensor(paddle::experimental::empty(
       {num_recv_tokens, hidden}, x.dtype(), x.place()));
@@ -905,8 +917,6 @@ Buffer::intranode_combine(
                      recv_topk_weights_ptr,
                      x.data_ptr(),
                      topk_weights_ptr,
-                     nullptr,  // bias_ptrs[0] (not exposed)
-                     nullptr,  // bias_ptrs[1] (not exposed)
                      src_idx.data_ptr<int>(),
                      rank_prefix_matrix.data_ptr<int>(),
                      channel_prefix_matrix.data_ptr<int>(),
@@ -1096,7 +1106,7 @@ Buffer::internode_dispatch(
 
   // FP8 scales checks
   float* x_scales_ptr = nullptr;
-  int num_scales = 0, scale_token_stride = 0, scale_hidden_stride = 0;
+  int num_scales = 0;
   if (x_scales.has_value()) {
     EP_HOST_ASSERT(x.element_size() == 1);
     EP_HOST_ASSERT(x_scales->scalar_type() == deep_ep::detail::kFloat32);
@@ -1105,8 +1115,6 @@ Buffer::internode_dispatch(
     EP_HOST_ASSERT(x_scales->size(0) == num_tokens);
     num_scales = x_scales->dim() == 1 ? 1 : static_cast<int>(x_scales->size(1));
     x_scales_ptr = x_scales->data_ptr<float>();
-    scale_token_stride = static_cast<int>(x_scales->stride(0));
-    scale_hidden_stride = static_cast<int>(x_scales->stride(1));
   }
 
   // Allocate all tensors on comm stream if set
@@ -1161,13 +1169,15 @@ Buffer::internode_dispatch(
         config.num_max_rdma_chunked_recv_tokens,
         buffer_ptrs_gpu,
         config.num_max_nvl_chunked_recv_tokens,
-        barrier_signal_ptrs_gpu,
+        task_fifo_ptrs_gpu,
+        head,
         rank,
         comm_stream,
         config.get_rdma_buffer_size_hint(hidden_int4 * sizeof(int4), num_ranks),
         num_nvl_bytes,
         true,
         low_latency_mode);
+    move_fifo_slots(2);
   } else {
     rdma_channel_prefix_matrix = ConvertPaddleTensorToDetailTensor(
         paddle::experimental::empty({num_rdma_ranks, num_channels},
@@ -1211,12 +1221,14 @@ Buffer::internode_dispatch(
         config.num_max_rdma_chunked_recv_tokens,
         buffer_ptrs_gpu,
         config.num_max_nvl_chunked_recv_tokens,
-        barrier_signal_ptrs_gpu,
+        task_fifo_ptrs_gpu,
+        head,
         rank,
         comm_stream,
         config.get_rdma_buffer_size_hint(hidden_int4 * sizeof(int4), num_ranks),
         num_nvl_bytes,
         low_latency_mode);
+    move_fifo_slots(3);
 
     // Synchronize total received tokens and tokens per expert
     auto start_time = std::chrono::high_resolution_clock::now();
@@ -1333,14 +1345,12 @@ Buffer::internode_dispatch(
       recv_rdma_rank_prefix_sum.data_ptr<int>(),
       gbl_channel_prefix_matrix.data_ptr<int>(),
       recv_gbl_rank_prefix_sum.data_ptr<int>(),
-      is_token_in_rank.data_ptr<bool>(),
       num_tokens,
       hidden_int4,
       num_scales,
       num_topk,
       num_experts,
-      scale_token_stride,
-      scale_hidden_stride,
+      is_token_in_rank.data_ptr<bool>(),
       rdma_buffer_ptr,
       config.num_max_rdma_chunked_send_tokens,
       config.num_max_rdma_chunked_recv_tokens,
@@ -1538,13 +1548,15 @@ Buffer::internode_combine(
       config.num_max_rdma_chunked_recv_tokens,
       buffer_ptrs_gpu,
       config.num_max_nvl_chunked_recv_tokens,
-      barrier_signal_ptrs_gpu,
+      task_fifo_ptrs_gpu,
+      head,
       rank,
       comm_stream,
       config.get_rdma_buffer_size_hint(hidden_int4 * sizeof(int4), num_ranks),
       num_nvl_bytes,
       false,
       low_latency_mode);
+  move_fifo_slots(2);
 
   // Launch data combine
   auto combined_x =
@@ -1556,8 +1568,6 @@ Buffer::internode_combine(
                      is_combined_token_in_rank.data_ptr<bool>(),
                      x.data_ptr(),
                      topk_weights_ptr,
-                     nullptr,  // bias_ptrs[0] (not exposed)
-                     nullptr,  // bias_ptrs[1] (not exposed)
                      combined_rdma_head.data_ptr<int>(),
                      combined_nvl_head.data_ptr<int>(),
                      src_meta.data_ptr(),

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

@@ -81,9 +81,10 @@ struct Buffer {
   // After IPC/NVSHMEM synchronization, this flag will be true
   bool available = false;
 
-  // Barrier signals
-  int* barrier_signal_ptrs[NUM_MAX_NVL_PEERS] = {nullptr};
-  int** barrier_signal_ptrs_gpu = nullptr;
+  // Task fifo
+  int head = 0;
+  int* task_fifo_ptrs[NUM_MAX_NVL_PEERS] = {nullptr};
+  int** task_fifo_ptrs_gpu = nullptr;
 
   // Workspace
   void* workspace = nullptr;
@@ -100,6 +101,9 @@ struct Buffer {
   volatile int* moe_recv_rdma_counter = nullptr;
   int* moe_recv_rdma_counter_mapped = nullptr;
 
+ private:
+  void move_fifo_slots(int num_slots = 1);
+
  public:
   Buffer(int rank,
          int num_ranks,

paddle/fluid/distributed/collective/deep_ep/include/types.h

@@ -73,8 +73,6 @@ struct Tensor {
   }
 
   int64_t element_size() const { return phi::SizeOf(raw_tensor_.dtype()); }
-
-  int64_t stride(int64_t d) const { return raw_tensor_.strides().at(d); }
 };
 
 }  // namespace deep_ep::detail