Optimize low latency combine recv kernel (about 3.0x speedup)

csrc/deep_ep.cpp

@@ -1170,7 +1170,11 @@ Buffer::low_latency_combine(const torch::Tensor& x, const torch::Tensor& topk_id
     EP_HOST_ASSERT(x.size(2) % sizeof(int4) == 0 and x.size(2) % 128 == 0);
     EP_HOST_ASSERT(topk_idx.dim() == 2 and topk_idx.is_contiguous());
     EP_HOST_ASSERT(topk_idx.size(0) == topk_weights.size(0) and topk_idx.size(1) == topk_weights.size(1));
-    EP_HOST_ASSERT(topk_idx.scalar_type() == torch::kInt64);
+
+    // HACK
+//    EP_HOST_ASSERT(topk_idx.scalar_type() == torch::kInt64);
+    EP_HOST_ASSERT(topk_idx.scalar_type() == torch::kInt32);
+
     EP_HOST_ASSERT(topk_weights.dim() == 2 and topk_weights.is_contiguous());
     EP_HOST_ASSERT(topk_weights.size(0) <= num_max_dispatch_tokens_per_rank);
     EP_HOST_ASSERT(topk_weights.scalar_type() == torch::kFloat32);
@@ -1214,7 +1218,10 @@ Buffer::low_latency_combine(const torch::Tensor& x, const torch::Tensor& topk_id
         internode_ll::combine(combined_x.data_ptr(),
                               buffer.combine_rdma_recv_data_buffer, buffer.combine_rdma_recv_flag_buffer,
                               buffer.combine_rdma_send_buffer,
-                              x.data_ptr(), topk_idx.data_ptr<int64_t>(), topk_weights.data_ptr<float>(),
+                              x.data_ptr(),
+//                              topk_idx.data_ptr<int64_t>(),
+                              topk_idx.data_ptr<int32_t>(),
+                              topk_weights.data_ptr<float>(),
                               src_info.data_ptr<int>(), layout_range.data_ptr<int64_t>(),
                               next_clean_meta.first, next_clean_meta.second,
                               num_combined_tokens, hidden, num_max_dispatch_tokens_per_rank,

csrc/kernels/api.cuh

@@ -152,7 +152,7 @@ void dispatch(void* packed_recv_x, void* packed_recv_x_scales,
 
 void combine(void* combined_x,
              void* rdma_recv_x, int* rdma_recv_flag, void* rdma_send_x,
-             const void* x, const int64_t* topk_idx, const float* topk_weights,
+             const void* x, const int32_t* topk_idx_i32, const float* topk_weights,
              const int* src_info, const int64_t* layout_range,
              int* next_clean, int num_next_clean_int,
              int num_combined_tokens, int hidden, int num_max_dispatch_tokens_per_rank,

csrc/kernels/internode_ll.cu

@@ -382,11 +382,18 @@ LAUNCH_KERNEL(&cfg, dispatch_func, \
 #undef DISPATCH_LAUNCH_CASE
 }
 
+// TODO generalize
+constexpr int kMaxNumTokensPerSm = 6;
+constexpr int kIdxOrWeightDim = 2;
+constexpr int kNumActualTopkDivFour = 2;
+constexpr int kNumActualTopk = kNumActualTopkDivFour * 4;
+constexpr int kWarpSize = 32;
+
 template <int kHidden, int kNumMaxTopk>
 __global__ __launch_bounds__(1024, 1) void
 combine(void* combined_x,
         void* rdma_recv_x, int* rdma_recv_flag, void* rdma_send_x,
-        const void* x, const int64_t* topk_idx, const float* topk_weights,
+        const void* x, const int32_t* topk_idx_i32, const float* topk_weights,
         const int* src_info, const int64_t* layout_range,
         int* next_clean, int num_next_clean_int,
         int* atomic_clean_flag,
@@ -489,42 +496,92 @@ combine(void* combined_x,
     if ((phases & LOW_LATENCY_RECV_PHASE) == 0)
         return;
 
+    int self_num_iteration = (sm_id >= num_combined_tokens) ? 0 : (1 + (num_combined_tokens - sm_id - 1) / num_sms);
+
+    alignas(16) __shared__ int4 shared_topk_info[kMaxNumTokensPerSm * kIdxOrWeightDim * kNumActualTopkDivFour];
+    const auto compute_shared_topk_info_addr = [=](int idx_iteration, int idx_iow, int idx_topkdivfour) {
+        return shared_topk_info
+            + idx_iteration * (kIdxOrWeightDim * kNumActualTopkDivFour)
+            + idx_iow * kNumActualTopkDivFour
+            + idx_topkdivfour;
+    };
+
+    int4 temp_buf;
+    int prepare_topk_idx_iteration, prepare_topk_idx_iow, prepare_topk_idx_topkdivfour;
+    static_assert(sizeof(shared_topk_info) / sizeof(shared_topk_info[0]) <= kWarpSize);
+    if (warp_id == 0) {
+        int index = thread_id;
+        prepare_topk_idx_topkdivfour = index % kNumActualTopkDivFour;
+        index /= kNumActualTopkDivFour;
+        prepare_topk_idx_iow = index % kIdxOrWeightDim;
+        index /= kIdxOrWeightDim;
+        prepare_topk_idx_iteration = index;
+    }
+    bool enable_prepare_topk = (warp_id == 0) and (prepare_topk_idx_iteration < self_num_iteration);
+    if (enable_prepare_topk) {
+        const int prepare_topk_token_idx = sm_id + prepare_topk_idx_iteration * num_sms;
+        const int4* src_addr = (
+            ((prepare_topk_idx_iow == 0)
+                ? reinterpret_cast<const int4*>(topk_idx_i32)
+                : reinterpret_cast<const int4*>(topk_weights))
+            + prepare_topk_token_idx * kNumActualTopkDivFour
+            + prepare_topk_idx_topkdivfour
+        );
+        temp_buf = ld_nc_global(src_addr);
+    }
+
     // Wait all ranks to arrive
-    if (responsible_expert_idx < num_experts) {
-        EP_DEVICE_ASSERT(num_warps_per_group > 1);
-        if (sub_warp_id == 0 and lane_id == 0) {
-            while (ld_acquire_sys_global(rdma_recv_flag + responsible_expert_idx) == 0);
-        }
+    const int recv_flag_responsible_expert_idx = thread_id;
+    if (recv_flag_responsible_expert_idx < num_experts) {
+        while (ld_acquire_sys_global(rdma_recv_flag + recv_flag_responsible_expert_idx) == 0);
+    }
+
+    if (enable_prepare_topk) {
+        int4* smem_addr = compute_shared_topk_info_addr(prepare_topk_idx_iteration, prepare_topk_idx_iow, prepare_topk_idx_topkdivfour);
+        *smem_addr = temp_buf;
     }
-    cg::this_grid().sync();
+
+    __syncthreads();
 
     // Reduce tokens with FP8 cast
     EP_DEVICE_ASSERT(num_topk <= 32 and hidden_bf16_int4 <= num_threads);
     EP_STATIC_ASSERT(kHidden % (32 * kNumElemsPerInt4) == 0, "Invalid vectorization");
     if (thread_id < hidden_bf16_int4) {
-        for (int token_idx = sm_id; token_idx < num_combined_tokens; token_idx += num_sms) {
+        for (int idx_iteration = 0; idx_iteration < self_num_iteration; ++ idx_iteration) {
+            const int token_idx = sm_id + idx_iteration * num_sms;
+
             // Read top-k indices and weights
-            int reg_topk_idx[kNumMaxTopk];
-            float reg_topk_weights[kNumMaxTopk];
+            alignas(16) int reg_topk_idx[kNumMaxTopk];
+            alignas(16) float reg_topk_weights[kNumMaxTopk];
+            auto reg_topk_idx_vec = reinterpret_cast<int4*>(reg_topk_idx);
+            auto reg_topk_weights_vec = reinterpret_cast<float4*>(reg_topk_weights);
+            #pragma unroll
+            for (int i = 0; i < kNumActualTopkDivFour; ++i) {
+                reg_topk_idx_vec[i] = *compute_shared_topk_info_addr(idx_iteration, 0, i);
+            }
             #pragma unroll
-            for (int i = 0; i < num_topk; ++ i) {
-                reg_topk_idx[i] = static_cast<int>(__ldg(topk_idx + token_idx * num_topk + i));
-                reg_topk_weights[i] = __ldg(topk_weights + token_idx * num_topk + i);
+            for (int i = 0; i < kNumActualTopkDivFour; ++i) {
+                reg_topk_weights_vec[i] = *reinterpret_cast<float4*>(compute_shared_topk_info_addr(idx_iteration, 1, i));
             }
 
-            float combined_values[kNumElemsPerInt4] = {0.0f};
+            // Read from sources, Reduce
+            int4 zero4 = {0,0,0,0};
+            int4 x_vec[kNumActualTopk];
             #pragma unroll
-            for (int i = 0; i < num_topk; ++ i) if (reg_topk_idx[i] >= 0) {
-                // Read from sources
+            for (int i = 0; i < kNumActualTopk; ++i) {
                 auto rdma_buffer_type = reinterpret_cast<const int*>(static_cast<uint8_t*>(rdma_recv_x) + (reg_topk_idx[i] * num_max_dispatch_tokens_per_rank + token_idx) * num_bytes_per_slot);
                 auto rdma_buffer_row = reinterpret_cast<const uint8_t*>(rdma_buffer_type);
+                x_vec[i] = (reg_topk_idx[i] >= 0) ? ld_nc_global(reinterpret_cast<const int4*>(rdma_buffer_row) + thread_id) : zero4;
+            }
 
-                // Reduce
-                auto x_vec = ld_nc_global(reinterpret_cast<const int4*>(rdma_buffer_row) + thread_id);
-                const auto x_bf16 = reinterpret_cast<nv_bfloat16*>(&x_vec);
-                #pragma unroll
-                for (int j = 0; j < kNumElemsPerInt4; ++ j)
-                    combined_values[j] += static_cast<float>(x_bf16[j]) * reg_topk_weights[i];
+            float combined_values[kNumElemsPerInt4] = {0.0f};
+            #pragma unroll
+            for (int i = 0; i < kNumActualTopk; ++i) {
+                if (reg_topk_idx[i] >= 0) {
+                    const auto x_bf16 = reinterpret_cast<nv_bfloat16*>(&x_vec[i]);
+                    #pragma unroll
+                    for (int j = 0; j < kNumElemsPerInt4; ++ j) combined_values[j] += static_cast<float>(x_bf16[j]) * reg_topk_weights[i];
+                }
             }
 
             // Write results
@@ -540,7 +597,7 @@ combine(void* combined_x,
 
 void combine(void* combined_x,
              void* rdma_recv_x, int* rdma_recv_flag, void* rdma_send_x,
-             const void* x, const int64_t* topk_idx, const float* topk_weights,
+             const void* x, const int32_t* topk_idx_i32, const float* topk_weights,
              const int* src_info, const int64_t* layout_range,
              int* next_clean, int num_next_clean_int,
              int num_combined_tokens, int hidden, int num_max_dispatch_tokens_per_rank,
@@ -559,13 +616,14 @@ void combine(void* combined_x,
     auto atomic_clean_flag = static_cast<int*>(workspace);
     EP_HOST_ASSERT(sizeof(int) <= NUM_WORKSPACE_BYTES);
     EP_HOST_ASSERT(num_topk <= kNumMaxTopk);
+    EP_HOST_ASSERT(num_combined_tokens <= kMaxNumTokensPerSm * num_sms);
 
 #define COMBINE_LAUNCH_CASE(hidden) { \
 auto combine_func = combine<hidden, kNumMaxTopk>; \
 LAUNCH_KERNEL(&cfg, combine_func, \
               combined_x, \
               rdma_recv_x, rdma_recv_flag, rdma_send_x, \
-              x, topk_idx, topk_weights, src_info, layout_range, \
+              x, topk_idx_i32, topk_weights, src_info, layout_range, \
               next_clean, num_next_clean_int, \
               atomic_clean_flag, \
               num_combined_tokens, hidden, num_topk, \

csrc/kernels/utils.cuh

@@ -199,6 +199,14 @@ __device__  __forceinline__ int4 ld_nc_global(const int4 *ptr) {
     return ret;
 }
 
+template <>
+__device__  __forceinline__ float4 ld_nc_global(const float4 *ptr) {
+    float4 ret;
+    asm volatile(LD_NC_FUNC ".v4.f32 {%0, %1, %2, %3}, [%4];"
+            : "=f"(ret.x), "=f"(ret.y), "=f"(ret.z), "=f"(ret.w) : "l"(ptr));
+    return ret;
+}
+
 __device__ __forceinline__ void st_na_relaxed(const uint8_t *ptr, uint8_t val) {
     asm volatile("st.relaxed.gpu.global.L1::no_allocate.b8 [%0], %1;" : : "l"(ptr), "h"(static_cast<uint16_t>(val)));
 }

deep_ep/buffer.py

@@ -570,6 +570,10 @@ def low_latency_combine(self, x: torch.Tensor, topk_idx: torch.Tensor, topk_weig
             event: the event after executing the kernel (valid only if `async_finish` is set).
             hook: the receiving hook function (valid only if `return_recv_hook` is set).
         """
+
+        # hack
+        topk_idx = topk_idx.to(torch.int32)
+
         src_info, layout_range, num_max_dispatch_tokens_per_rank, hidden, num_experts = handle
         combined_x, event, hook = self.runtime.low_latency_combine(x, topk_idx, topk_weights, src_info, layout_range,
                                                                    num_max_dispatch_tokens_per_rank, num_experts,

setup.py

@@ -21,6 +21,9 @@
     nvcc_dlink = []
     extra_link_args = []
 
+    # TODO merge another PR first
+    nvcc_flags += ['-lineinfo']
+
     # NVSHMEM flags
     if disable_nvshmem:
         cxx_flags.append('-DDISABLE_NVSHMEM')

tests/test_internode.py

@@ -248,5 +248,5 @@ def test_loop(local_rank: int, num_local_ranks: int):
 
 
 if __name__ == '__main__':
-    num_processes = 8
+    num_processes = int(os.getenv("DEEPEP_TEST_NUM_PROCESSES", "8"))
     torch.multiprocessing.spawn(test_loop, args=(num_processes, ), nprocs=num_processes)

tests/test_intranode.py

@@ -1,3 +1,4 @@
+import os
 import time
 import torch
 import torch.distributed as dist
@@ -12,7 +13,12 @@
 
 def test_main(num_sms: int, local_rank: int, num_ranks: int, rank: int, buffer: deep_ep.Buffer, group: dist.ProcessGroup):
     # Settings
-    num_tokens, hidden, num_topk, num_experts = 4096, 7168, 8, (256 // num_ranks) * num_ranks
+    # num_tokens, hidden, num_topk, num_experts = 4096, 7168, 8, (256 // num_ranks) * num_ranks
+    num_tokens = int(os.environ.get("DEEPEP_TEST_NUM_TOKENS", "4096"))
+    hidden = int(os.environ.get("DEEPEP_TEST_HIDDEN", "7168"))
+    num_topk = int(os.environ.get("DEEPEP_TEST_NUM_TOPK", "8"))
+    num_experts = int(os.environ.get("DEEPEP_TEST_NUM_EXPERTS", str((256 // num_ranks) * num_ranks)))
+
     assert num_experts % num_ranks == 0
     if local_rank == 0:
         print(f'[config] num_tokens={num_tokens}, hidden={hidden}, num_topk={num_topk}', flush=True)
@@ -184,9 +190,9 @@ def check_data(check_x, rank_prefix_matrix):
                 best_time, best_results = t, (num_sms, nvl_chunk_size)
             if local_rank == 0:
                 print(f'[tuning] SMs {num_sms}, NVL chunk {nvl_chunk_size if nvl_chunk_size else "default"}: '
-                      f'{nvl_recv_bytes / 1e9 / t:.2f} GB/s (NVL) ', flush=True)
+                      f'{nvl_recv_bytes / 1e9 / t:.2f} GB/s (NVL) t={t * 1e3}ms', flush=True)
         if local_rank == 0:
-            print(f'[tuning] Best dispatch ({"FP8" if isinstance(current_x, tuple) else "BF16"}): SMs {best_results[0]}, NVL chunk {best_results[1]}, {nvl_recv_bytes / 1e9 / best_time:.2f} GB/s (NVL)', flush=True)
+            print(f'[tuning] Best dispatch ({"FP8" if isinstance(current_x, tuple) else "BF16"}): SMs {best_results[0]}, NVL chunk {best_results[1]}, {nvl_recv_bytes / 1e9 / best_time:.2f} GB/s (NVL) t={best_time * 1e3}ms', flush=True)
             print('', flush=True)
 
         # Gather the best config from rank 0 and the first test setting
@@ -215,12 +221,12 @@ def check_data(check_x, rank_prefix_matrix):
         t = bench(lambda: buffer.combine(**tune_args))[0]
         if local_rank == 0:
             print(f'[tuning] SMs {num_sms}, NVL chunk {nvl_chunk_size if nvl_chunk_size else "default"}: '
-                  f'{combine_bf16_nvl_send_bytes / 1e9 / t:.2f} GB/s (NVL) ', flush=True)
+                  f'{combine_bf16_nvl_send_bytes / 1e9 / t:.2f} GB/s (NVL)  t={t * 1e3}ms', flush=True)
             if t < best_time and nvl_chunk_size > 0:
                 best_time, best_results = t, (num_sms, nvl_chunk_size)
 
     if local_rank == 0:
-        print(f'[tuning] Best combine: SMs {best_results[0]}, NVL chunk {best_results[1]}: {combine_bf16_nvl_send_bytes / 1e9 / best_time:.2f} GB/s (NVL)', flush=True)
+        print(f'[tuning] Best combine: SMs {best_results[0]}, NVL chunk {best_results[1]}: {combine_bf16_nvl_send_bytes / 1e9 / best_time:.2f} GB/s (NVL) t={best_time * 1e3}ms', flush=True)
         print('', flush=True)
 
 
@@ -236,7 +242,9 @@ def test_loop(local_rank: int, num_local_ranks: int):
                             num_qps_per_rank=(ll_num_experts // num_ranks if test_ll_compatibility else 1))
     torch.manual_seed(rank)
 
-    for i in (24, ):
+    num_sms = int(os.environ.get("DEEPEP_TEST_NUM_SMS", "24"))
+
+    for i in (num_sms, ):
         test_main(i, local_rank, num_ranks, rank, buffer, group)
         if local_rank == 0:
             print('', flush=True)
@@ -252,5 +260,5 @@ def test_loop(local_rank: int, num_local_ranks: int):
 
 
 if __name__ == '__main__':
-    num_processes = 8
+    num_processes = int(os.getenv("DEEPEP_TEST_NUM_PROCESSES", "8"))
     torch.multiprocessing.spawn(test_loop, args=(num_processes, ), nprocs=num_processes)