commit

Author: zhoutianzi666

paddle/fluid/distributed/collective/deep_ep/kernels/internode_ll.cu

@@ -658,13 +658,13 @@ void dispatch(void* packed_recv_x,
               cudaStream_t stream,
               int phases,
               int num_per_channel) {
-  constexpr int kNumMaxTopK = 9;
+  constexpr int kNumMaxTopK = 16;
   constexpr int NUM_WARPS = 32;
   const int dev_id = 0;
   int sm_count;
   cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, dev_id);
   const int num_warp_groups = cell_div(num_experts, sm_count);
-  const auto num_sms = max(sm_count, cell_div(num_experts, num_warp_groups));
+
   EP_HOST_ASSERT(num_topk <= kNumMaxTopK);
 
   // Workspace checks
@@ -681,6 +681,7 @@ void dispatch(void* packed_recv_x,
           kNumPerChannels,
           {DISPATCH_NUM_WARP_GROUPS(num_warp_groups, kNumWarpGroups, {     // 1
             constexpr int kNumWarpsPerGroup = NUM_WARPS / kNumWarpGroups;  // 32
+            // because of many `warp_id < num_topk` in kernel.
             EP_STATIC_ASSERT(
                 kNumMaxTopK + 1 <= kNumWarpGroups * kNumWarpsPerGroup,
                 "Too many top-k selections");
@@ -695,7 +696,7 @@ void dispatch(void* packed_recv_x,
                                                     kHidden,
                                                     kNumPerChannels>;
             SETUP_LAUNCH_CONFIG(
-                num_sms, kNumWarpGroups * kNumWarpsPerGroup * 32, stream);
+                sm_count, kNumWarpGroups * kNumWarpsPerGroup * 32, stream);
             LAUNCH_KERNEL(&cfg,
                           dispatch_func,
                           packed_recv_x,
@@ -993,14 +994,13 @@ void combine(void* combined_x,
              cudaStream_t stream,
              int phases,
              bool zero_copy) {
-  constexpr int kNumMaxTopk = 9;
+  constexpr int kNumMaxTopk = 16;
   constexpr int NUM_WARPS = 32;
 
   const int dev_id = 0;
   int sm_count;
   cudaDeviceGetAttribute(&sm_count, cudaDevAttrMultiProcessorCount, dev_id);
   const int num_warp_groups = cell_div(num_experts, sm_count);
-  const auto num_sms = max(sm_count, cell_div(num_experts, num_warp_groups));
 
   // Check workspace
   auto atomic_clean_flag = reinterpret_cast<int*>(workspace);
@@ -1015,7 +1015,7 @@ void combine(void* combined_x,
         auto combine_func =
             combine<kNumWarpGroups, kNumWarpsPerGroup, kHidden, kNumMaxTopk>;
         SETUP_LAUNCH_CONFIG(
-            num_sms, kNumWarpGroups * kNumWarpsPerGroup * 32, stream);
+            sm_count, kNumWarpGroups * kNumWarpsPerGroup * 32, stream);
         LAUNCH_KERNEL(&cfg,
                       combine_func,
                       combined_x,