Update deep_ep intranode & internode kernels

Author: lshpku

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

@@ -69,14 +69,15 @@ Buffer::Buffer(int rank,
   calc_ctx = reinterpret_cast<phi::GPUContext*>(
       reinterpret_cast<paddle::distributed::ProcessGroupNCCL*>(pg)
           ->GetDeviceContext(place, true));
-  // 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;
+
+  // 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*);
 
   // Common checks
   EP_HOST_ASSERT(num_nvl_bytes % NUM_BUFFER_ALIGNMENT_BYTES == 0 &&
-                 (num_nvl_bytes <= std::numeric_limits<int64_t>::max() ||
+                 (num_nvl_bytes <= std::numeric_limits<int>::max() ||
                   num_rdma_bytes == 0));
   EP_HOST_ASSERT(
       num_rdma_bytes % NUM_BUFFER_ALIGNMENT_BYTES == 0 &&
@@ -90,40 +91,35 @@ 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 handle
-    CUDA_CHECK(cudaMalloc(
-        &buffer_ptrs[nvl_rank],
-        num_nvl_bytes + fifo_bytes + buffer_ptr_bytes + task_ptr_bytes));
+    // 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));
     CUDA_CHECK(
         cudaIpcGetMemHandle(&ipc_handles[nvl_rank], buffer_ptrs[nvl_rank]));
-    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);
+    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);
 
     // No need to synchronize, will do a full device sync during `sync`
     CUDA_CHECK(cudaMemsetAsync(
-        buffer_ptrs[nvl_rank],
-        0,
-        num_nvl_bytes + fifo_bytes + buffer_ptr_bytes + task_ptr_bytes,
-        comm_stream));
+        barrier_signal_ptrs[nvl_rank], 0, barrier_signal_bytes, comm_stream));
   }
 
   // Create 32 MiB workspace
@@ -165,8 +161,7 @@ Buffer::~Buffer() noexcept(false) {
   if (num_nvl_bytes > 0) {
     // Barrier
     intranode::barrier(
-        task_fifo_ptrs_gpu, head, nvl_rank, num_nvl_ranks, comm_stream);
-    move_fifo_slots();
+        barrier_signal_ptrs_gpu, nvl_rank, num_nvl_ranks, comm_stream);
     CUDA_CHECK(cudaDeviceSynchronize());
 
     // Close remote IPC
@@ -197,10 +192,6 @@ 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 {
@@ -249,7 +240,7 @@ void Buffer::sync(
 
   // Sync IPC handles
   if (num_nvl_bytes > 0) {
-    EP_HOST_ASSERT(num_ranks == static_cast<int64_t>(device_ids.size()));
+    EP_HOST_ASSERT(num_ranks == 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) {
@@ -261,22 +252,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));
-        task_fifo_ptrs[i] = reinterpret_cast<int*>(
-            reinterpret_cast<uint8_t*>(buffer_ptrs[i]) + num_nvl_bytes);
+        barrier_signal_ptrs[i] = reinterpret_cast<int*>(
+            static_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 task pointers to GPU
+    // Copy all buffer and barrier signal pointers to GPU
     CUDA_CHECK(cudaMemcpy(buffer_ptrs_gpu,
                           buffer_ptrs,
                           sizeof(void*) * NUM_MAX_NVL_PEERS,
                           cudaMemcpyHostToDevice));
-    CUDA_CHECK(cudaMemcpy(task_fifo_ptrs_gpu,
-                          task_fifo_ptrs,
+    CUDA_CHECK(cudaMemcpy(barrier_signal_ptrs_gpu,
+                          barrier_signal_ptrs,
                           sizeof(int*) * NUM_MAX_NVL_PEERS,
                           cudaMemcpyHostToDevice));
     CUDA_CHECK(cudaDeviceSynchronize());
@@ -520,7 +511,7 @@ Buffer::intranode_dispatch(
 
   // FP8 scales checks
   float* x_scales_ptr = nullptr;
-  int num_scales = 0;
+  int num_scales = 0, scale_token_stride = 0, scale_hidden_stride = 0;
   if (x_scales.has_value()) {
     EP_HOST_ASSERT(x.element_size() == 1);
     EP_HOST_ASSERT(x_scales->scalar_type() == deep_ep::detail::kFloat32);
@@ -529,6 +520,8 @@ 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
@@ -564,12 +557,10 @@ Buffer::intranode_dispatch(
     intranode::cached_notify_dispatch(rank_prefix_matrix.data_ptr<int>(),
                                       num_memset_int,
                                       buffer_ptrs_gpu,
-                                      task_fifo_ptrs_gpu,
-                                      head,
+                                      barrier_signal_ptrs_gpu,
                                       rank,
                                       num_ranks,
                                       comm_stream);
-    move_fifo_slots(2);
   } else {
     rank_prefix_matrix = ConvertPaddleTensorToDetailTensor(
         paddle::experimental::empty({num_ranks, num_ranks},
@@ -604,12 +595,10 @@ Buffer::intranode_dispatch(
                                num_memset_int,
                                expert_alignment,
                                buffer_ptrs_gpu,
-                               task_fifo_ptrs_gpu,
-                               head,
+                               barrier_signal_ptrs_gpu,
                                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();
@@ -719,10 +708,13 @@ 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,
@@ -867,15 +859,11 @@ Buffer::intranode_combine(
                                    num_channels,
                                    num_recv_tokens,
                                    num_channels * num_ranks * 2,
-                                   task_fifo_ptrs_gpu,
-                                   head,
+                                   barrier_signal_ptrs_gpu,
                                    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()));
@@ -895,6 +883,8 @@ 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>(),
@@ -1084,7 +1074,7 @@ Buffer::internode_dispatch(
 
   // FP8 scales checks
   float* x_scales_ptr = nullptr;
-  int num_scales = 0;
+  int num_scales = 0, scale_token_stride = 0, scale_hidden_stride = 0;
   if (x_scales.has_value()) {
     EP_HOST_ASSERT(x.element_size() == 1);
     EP_HOST_ASSERT(x_scales->scalar_type() == deep_ep::detail::kFloat32);
@@ -1093,6 +1083,8 @@ 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
@@ -1144,15 +1136,13 @@ Buffer::internode_dispatch(
         config.num_max_rdma_chunked_recv_tokens,
         buffer_ptrs_gpu,
         config.num_max_nvl_chunked_recv_tokens,
-        task_fifo_ptrs_gpu,
-        head,
+        barrier_signal_ptrs_gpu,
         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},
@@ -1196,14 +1186,12 @@ Buffer::internode_dispatch(
         config.num_max_rdma_chunked_recv_tokens,
         buffer_ptrs_gpu,
         config.num_max_nvl_chunked_recv_tokens,
-        task_fifo_ptrs_gpu,
-        head,
+        barrier_signal_ptrs_gpu,
         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();
@@ -1320,12 +1308,14 @@ 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,
-      is_token_in_rank.data_ptr<bool>(),
+      scale_token_stride,
+      scale_hidden_stride,
       rdma_buffer_ptr,
       config.num_max_rdma_chunked_send_tokens,
       config.num_max_rdma_chunked_recv_tokens,
@@ -1523,15 +1513,13 @@ Buffer::internode_combine(
       config.num_max_rdma_chunked_recv_tokens,
       buffer_ptrs_gpu,
       config.num_max_nvl_chunked_recv_tokens,
-      task_fifo_ptrs_gpu,
-      head,
+      barrier_signal_ptrs_gpu,
       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 =
@@ -1543,6 +1531,8 @@ 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

@@ -77,10 +77,9 @@ struct Buffer {
   // After IPC/NVSHMEM synchronization, this flag will be true
   bool available = false;
 
-  // Task fifo
-  int head = 0;
-  int* task_fifo_ptrs[NUM_MAX_NVL_PEERS] = {nullptr};
-  int** task_fifo_ptrs_gpu = nullptr;
+  // Barrier signals
+  int* barrier_signal_ptrs[NUM_MAX_NVL_PEERS] = {nullptr};
+  int** barrier_signal_ptrs_gpu = nullptr;
 
   // Workspace
   void* workspace = nullptr;
@@ -97,9 +96,6 @@ 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,6 +73,8 @@ 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