分析一下EP并行和DeepSeek开源的DeepEP代码

1. EP并行概念
2. SGlang EP并行实现
3. DeepEP
3.1 用于通信的Buffer
3.2 用于训练和Prefill的高吞吐Kernel

3.2.1 Dispatch通信流程
3.2.1.1 Notify_Dispatch
3.2.1.2 Intranode::dispatch
3.2.1.3 Internode::dispatch
3.2.1.3.1 kRDMASender
3.2.1.3.2 kRDMASenderCoordinator
3.2.1.3.3  kRDMAAndNVLForwarder
3.2.1.3.4 kForwarderCoordinator
3.2.1.3.5 kNVLReceivers

3.2.2 Combine流程
3.2.2.1 Intranode_Combine
3.2.2.2 Internode_Combine
3.2.2.2.1 kNVLSender
3.2.2.2.2 kNVLAndRDMAForwarder
3.2.2.2.3 kRDMAReceiver
3.2.2.2.4 kCoordinator

3.2 用于Decoding的低延迟Kernel
3.2.1 LowLatency Layout

3.2.2 低延迟Dispatch
3.2.2.1 SEND PHASE
3.2.2.2 RECV PHASE

3.2.3 低延迟Combine
3.2.3.1 SEND PHASE
3.2.3.2 RECV PHASE

3.3 其它细节
3.3.1 文档外行为的PTX指令
3.3.2 Memory Order
3.3.3 nvshmem库的修改

4. RoCE上运行DeepEP的挑战
4.1 DeepSeek用到的IB网络技术
4.2 RoCE使用DeepEP的问题
4.2.1 Multi-Rail and Rail-Only拓扑的问题
4.2.2 incast
4.2.3 RC兼容
4.2.4 In Network Computing如何做?
4.3 真正的适合EP的RDMA over Ethernet

5. 关于DeepSeek-V3论文的建议
6. 关于MoE模型的演进

dim = 7168
inter_dim = 2048
tokens = 256
e = Expert(dim, inter_dim)

from ptflops import get_model_complexity_info
input_tokens = (1,tokens,dim)
flops, params = get_model_complexity_info(e, input_tokens, as_strings=True,print_per_layer_stat=True)

Expert(
  44.05 M, 100.000% Params, 11.28 GMac, 99.995% MACs, 
  (w1): Linear(14.68 M, 33.329% Params, 3.76 GMac, 33.328% MACs, in_features=7168, out_features=2048, bias=True)
  (w2): Linear(14.69 M, 33.341% Params, 3.76 GMac, 33.340% MACs, in_features=2048, out_features=7168, bias=True)
  (w3): Linear(14.68 M, 33.329% Params, 3.76 GMac, 33.328% MACs, in_features=7168, out_features=2048, bias=True)

class EPMoE(torch.nn.Module):
  def __init__():      
        self.num_experts = num_experts
        assert self.num_experts % self.tp_size == 0
        
        ## 根据节点的Rank计算出EP并行时, 该节点需要放置的专家.
        self.num_experts_per_partition = self.num_experts // self.tp_size
        self.start_expert_id = self.tp_rank * self.num_experts_per_partition
        self.end_expert_id = self.start_expert_id + self.num_experts_per_partition - 1

def forward(self, hidden_states: torch.Tensor, router_logits: torch.Tensor):
     ## 初始化GroupGemm
     
     ## 选择专家, 并根据输出的topk_ids进行重排序, 便于单个专家权重加载的时候复用内存处理多个token
      topk_weights, topk_ids = select_experts()
      reorder_topk_ids, src2dst, seg_indptr = run_moe_ep_preproess(
            topk_ids, self.num_experts
      )
      
      # 重排序数据的预处理过程
      pre_reorder_triton_kernel[(hidden_states.shape[0],)]()
      
      # 基于当前rank获取需要处理的segment
      seg_indptr_cur_rank = seg_indptr[self.start_expert_id : self.end_expert_id + 2]
      weight_indices_cur_rank = torch.arange(
            0,
            self.num_experts_per_partition,
            device=hidden_states.device,
            dtype=torch.int64,
        )   
        
        # 第一次GroupGemm做个Expert的Gate和Up-projection矩阵运算
        gateup_output = self.grouped_gemm_runner(
            a=gateup_input,
            b=self.w13_weight,
            c=gateup_output,
            ....)
            
        # Expert 的激活函数处理
         if self.activation == "silu":
            silu_and_mul_triton_kernel[(gateup_output.shape[0],)]()
            
        # Expert的 down projection GroupGEMM处理
         down_output = self.grouped_gemm_runner(
            a=down_input,
            b=self.w2_weight,
            c=down_output,)
        
        # Post重排
        post_reorder_triton_kernel[(hidden_states.size(0),)](
            down_output,
            output,
            src2dst,...)

from deep_ep import Buffer, EventOverlap

_buffer: Optional[Buffer] = None
def get_buffer(group: dist.ProcessGroup, hidden_bytes: int) -> Buffer:
    # 计算NVLink和RDMA需要多少Bytes
    num_nvl_bytes, num_rdma_bytes = 0, 0
    for config in (Buffer.get_dispatch_config(group.size()), Buffer.get_combine_config(group.size())):
        num_nvl_bytes = max(config.get_nvl_buffer_size_hint(hidden_bytes, group.size()), num_nvl_bytes)
        num_rdma_bytes = max(config.get_rdma_buffer_size_hint(hidden_bytes, group.size()), num_rdma_bytes)
    # 为Buffer分配内存
    if _buffer is None or _buffer.group != group or _buffer.num_nvl_bytes < num_nvl_bytes or _buffer.num_rdma_bytes < num_rdma_bytes:
        _buffer = Buffer(group, num_nvl_bytes, num_rdma_bytes)
    return _buffer

# Intranode
if num_ranks <= 8:
     return Config(Buffer.num_sms, 6, 256, 6, 128)

# Internode
 config_map = {
    16: Config(Buffer.num_sms, 16, 288, 20, 128),
    24: Config(Buffer.num_sms, 8, 288, 32, 128),
    32: Config(Buffer.num_sms, 8, 288, 32, 128),
    64: Config(Buffer.num_sms, 20, 288, 28, 128),
    128: Config(Buffer.num_sms, 20, 560, 32, 128),
    144: Config(Buffer.num_sms, 32, 720, 12, 128),
    160: Config(Buffer.num_sms, 28, 720, 12, 128),
  }

struct Config {
    int num_sms;
    int num_max_nvl_chunked_send_tokens;
    int num_max_nvl_chunked_recv_tokens;
    int num_max_rdma_chunked_send_tokens;
    int num_max_rdma_chunked_recv_tokens;
 }

// 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;

// Host-side MoE info
    volatile int* moe_recv_counter = nullptr;
    int* moe_recv_counter_mapped = nullptr;

    // Host-side expert-level MoE info
    volatile int* moe_recv_expert_counter = nullptr;
    int* moe_recv_expert_counter_mapped = nullptr;

    // Host-side RDMA-level MoE info
    volatile int* moe_recv_rdma_counter = nullptr;
    int* moe_recv_rdma_counter_mapped = nullptr;

# Calculate layout before actual dispatch
    num_tokens_per_rank, num_tokens_per_rdma_rank, num_tokens_per_expert, is_token_in_rank, previous_event = \
        _buffer.get_dispatch_layout(topk_idx, num_experts,
                                    previous_event=previous_event, async_finish=True,
                                    allocate_on_comm_stream=previous_event is not None)

Buffer::intranode_dispatch(
  const torch::Tensor& x,//Tensor数据x
conststd::optional<torch::Tensor>& x_scales, //用于FP8缩放的值
conststd::optional<torch::Tensor>& topk_idx,// MoE Gating产生的专家index
conststd::optional<torch::Tensor>& topk_weights, //MoE Gating产生的weight
conststd::optional<torch::Tensor>& num_tokens_per_rank, //每个Rank的token数
const torch::Tensor& is_token_in_rank, //标记token是否从某个rank发送
conststd::optional<torch::Tensor>& num_tokens_per_expert, //基于专家的token数统计
int cached_num_recv_tokens, //需要缓存的token数
conststd::optional<torch::Tensor>& cached_rank_prefix_matrix,  //[num_ranks,num_ranks]矩阵
conststd::optional<torch::Tensor>& cached_channel_prefix_matrix,//[num_ranks, num_channels]的矩阵
int expert_alignment, 
const Config& config, 
std::optional<EventHandle>& previous_event, 
bool async, bool allocate_on_comm_stream //是否需要异步执行和分配通信的cudastream
) {

auto compute_stream = at::cuda::getCurrentCUDAStream();
    if (allocate_on_comm_stream) {
        EP_HOST_ASSERT(previous_event.has_value() and async);
        at::cuda::setCurrentCUDAStream(comm_stream);
    }

    // Wait previous tasks to be finished
    if (previous_event.has_value()) {
        stream_wait(comm_stream, previous_event.value());
    } else {
        stream_wait(comm_stream, compute_stream);
    }

//初始化接收counter
*moe_recv_counter = -1, *moe_recv_rdma_counter = -1;
for (int i = 0; i < num_local_experts; ++ i)
    moe_recv_expert_counter[i] = -1;

//调用notify_dispath
internode::notify_dispatch(
  num_tokens_per_rank->data_ptr<int>(), //每个Rank需要发送的token数量
  moe_recv_counter_mapped, 
  num_ranks,
  num_tokens_per_rdma_rank->data_ptr<int>(),  //dispatch_layout的结果,在每个RDMA Rank上需要发送的token数量
  moe_recv_rdma_counter_mapped,
  num_tokens_per_expert->data_ptr<int>(), //dispatch_layout的结果,基于专家的token数统计
  moe_recv_expert_counter_mapped, 
  num_experts,
  is_token_in_rank.data_ptr<bool>(), //[num_tokens, num_ranks]二维数组, 数据类型为bool型, 标记token是否从某个rank发送.
  num_tokens, num_channels, //token数和Channel数
  hidden_int4, num_scales, num_topk, expert_alignment,
  rdma_channel_prefix_matrix.data_ptr<int>(), //{num_rdma_ranks, num_channels}
  recv_rdma_rank_prefix_sum.data_ptr<int>(), //{num_rdma_ranks}
  gbl_channel_prefix_matrix.data_ptr<int>(), //{num_ranks, num_channels}
  recv_gbl_rank_prefix_sum.data_ptr<int>(), //{num_ranks}
  rdma_buffer_ptr, 
  config.num_max_rdma_chunked_recv_tokens,
  buffer_ptrs_gpu, 
  config.num_max_nvl_chunked_recv_tokens,
  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);

// Allocate new tensors
    auto recv_x = torch::empty({num_recv_tokens, hidden}, x.options());
    auto recv_topk_idx = std::optional<torch::Tensor>(), recv_topk_weights = std::optional<torch::Tensor>(), recv_x_scales = std::optional<torch::Tensor>();
    auto recv_src_meta = std::optional<torch::Tensor>();
    auto recv_rdma_channel_prefix_matrix = std::optional<torch::Tensor>();
    auto recv_gbl_channel_prefix_matrix = std::optional<torch::Tensor>();
    auto send_rdma_head = std::optional<torch::Tensor>();
    auto send_nvl_head = std::optional<torch::Tensor>();
    if (not cached_mode) {
        recv_src_meta = torch::empty({num_recv_tokens, internode::get_source_meta_bytes()}, dtype(torch::kByte).device(torch::kCUDA));
        recv_rdma_channel_prefix_matrix = torch::empty({num_rdma_ranks, num_channels}, dtype(torch::kInt32).device(torch::kCUDA));
        recv_gbl_channel_prefix_matrix = torch::empty({num_ranks, num_channels}, dtype(torch::kInt32).device(torch::kCUDA));
        send_rdma_head = torch::empty({num_tokens, num_rdma_ranks}, dtype(torch::kInt32).device(torch::kCUDA));
        send_nvl_head = torch::empty({num_rdma_recv_tokens, NUM_MAX_NVL_PEERS}, dtype(torch::kInt32).device(torch::kCUDA));
    }

enum class WarpRole {
        kRDMASender, //负责将本地数据通过RDMA发送到远程节点
        kRDMASenderCoordinator, //管理RDMA发送的进度与同步
        kRDMAAndNVLForwarder, //将RDMA接收的数据转发到本地NVLink缓冲区
        kForwarderCoordinator, //全局协调转发任务
        kNVLReceivers //从NVLink缓冲区读取数据并写入最终接收数组
    };

for (int dst_rdma_rank = warp_id; dst_rdma_rank < kNumRDMARanks; dst_rdma_rank += kNumDispatchRDMASenderWarps) {
    // 根据不同的lane_id设置send_buffer的内容
    nvshmemx_int_put_nbi_warp(rdma_channel_meta.recv_buffer(rdma_rank), rdma_channel_meta.send_buffer(dst_rdma_rank), NUM_MAX_NVL_PEERS * 2 + 2,
                              translate_dst_rdma_rank<kLowLatencyMode>(dst_rdma_rank, nvl_rank));
}
nvshmem_fence();
sync_rdma_sender_smem();

// Copy `x` into symmetric send buffer
            auto st_broadcast = [=](constint key, const int4& value) {
                #pragma unroll
                for (int j = 0; j < num_topk_ranks; ++ j)
                    st_na_global(reinterpret_cast<int4*>(dst_send_buffers[j]) + key, value);
            };
            UNROLLED_WARP_COPY(5, lane_id, hidden_int4, 0, x + token_idx * hidden_int4, ld_nc_global, st_broadcast);
            #pragma unroll
            for (int i = 0; i < num_topk_ranks; ++ i)
                dst_send_buffers[i] = reinterpret_cast<int4*>(dst_send_buffers[i]) + hidden_int4;

            // Copy source metadata into symmetric send buffer
            if (lane_id < num_topk_ranks)
                st_na_global(reinterpret_cast<SourceMeta*>(dst_send_buffers[lane_id]), src_meta);
            #pragma unroll
            for (int i = 0; i < num_topk_ranks; ++ i)
                dst_send_buffers[i] = reinterpret_cast<SourceMeta*>(dst_send_buffers[i]) + 1;

            // Copy `x_scales` into symmetric send buffer
            #pragma unroll
            for (int i = lane_id; i < num_scales; i += 32) {
                auto value = ld_nc_global(x_scales + token_idx * num_scales + i);
                #pragma unroll
                for (int j = 0; j < num_topk_ranks; ++ j)
                    st_na_global(reinterpret_cast<float*>(dst_send_buffers[j]) + i, value);
            }
            #pragma unroll
            for (int i = 0; i < num_topk_ranks; ++ i)
                dst_send_buffers[i] = reinterpret_cast<float*>(dst_send_buffers[i]) + num_scales;

            // Copy `topk_idx` and `topk_weights` into symmetric send buffer
            #pragma unroll
            for (int i = lane_id; i < num_topk * num_topk_ranks; i += 32) {
                auto rank_idx = i / num_topk, copy_idx = i % num_topk;
                auto idx_value = static_cast<int>(ld_nc_global(topk_idx + token_idx * num_topk + copy_idx));
                auto weight_value = ld_nc_global(topk_weights + token_idx * num_topk + copy_idx);
                st_na_global(reinterpret_cast<int*>(dst_send_buffers[rank_idx]) + copy_idx, idx_value);
                st_na_global(reinterpret_cast<float*>(dst_send_buffers[rank_idx]) + num_topk + copy_idx, weight_value);
            }

// `st.global.L1::no_allocate` will be translated into `ST.E.NA.[width]` in SASS,
// which does not have cache allocation (obviously in L1, I guess not in L2 too)
#ifndef DISABLE_AGGRESSIVE_PTX_INSTRS
#define ST_NA_FUNC "st.global.L1::no_allocate"
#else
#define ST_NA_FUNC "st.global"
#endif

return recv_x, recv_topk_idx, recv_topk_weights, num_recv_tokens_per_expert_list, handle, event

def combine(self, x: torch.Tensor, handle: Tuple,
                topk_weights: Optional[torch.Tensor] = None,
                config: Optional[Config] = None,
                previous_event: Optional[EventOverlap] = None, async_finish: bool = False,
                allocate_on_comm_stream: bool = False) -> \
            Tuple[torch.Tensor, Optional[torch.Tensor], EventOverlap]:

# intranode handle
handle = (rank_prefix_matrix, 
     channel_prefix_matrix, 
     recv_channel_prefix_matrix, 
     recv_src_idx, is_token_in_rank, send_head)

# internode handle
handle = (is_token_in_rank,
    rdma_channel_prefix_matrix, gbl_channel_prefix_matrix,
    recv_rdma_channel_prefix_matrix, 
    recv_rdma_rank_prefix_sum, 
    recv_gbl_channel_prefix_matrix, 
    recv_gbl_rank_prefix_sum,
    recv_src_meta, send_rdma_head, send_nvl_head)

constexpr int kDtypePerInt4 = sizeof(int4) / sizeof(dtype_t);
    int hidden_int4 = hidden * sizeof(dtype_t) / sizeof(int4);
    auto x_int4 = reinterpret_cast<const int4*>(x);
    auto recv_int4 = reinterpret_cast<int4*>(recv_x);

// Shared head, tail and retired flags for receiver warps
        __shared__ volatile int warp_channel_head_idx[num_recv_warps][kNumRanks];//每个warp的channel头索引
        __shared__ volatile int channel_tail_idx[kNumRanks];  //每个Rank的channel尾索引
        __shared__ volatile bool warp_retired[num_recv_warps]; //每个接收 warp 是否已Retried的标志

// Reduce data
#pragma unroll
for (int i = recv_lane_id; i < hidden_int4; i += 32) {
     // Read buffers
     int4 recv_value_int4[kNumRanks];
     #pragma unroll
     for (int j = 0; j < num_topk_ranks; ++ j)
         recv_value_int4[j] = ld_nc_global(channel_x_buffers[topk_ranks[j]].buffer() + slot_indices[j] * hidden_int4 + i);
     // Reduce all-to-all results
     float values[kDtypePerInt4] = {0};
     #pragma unroll
     for (int j = 0; j < num_topk_ranks; ++ j) {
         auto recv_value_dtypes = reinterpret_cast<constdtype_t*>(&recv_value_int4[j]);
         #pragma unroll
         for (int k = 0; k < kDtypePerInt4; ++ k)
             values[k] += static_cast<float>(recv_value_dtypes[k]);
     }
     // Cast back to `dtype_t` and write
     int4 out_int4;
     auto out_dtypes = reinterpret_cast<dtype_t*>(&out_int4);
     #pragma unroll
     for (int j = 0; j < kDtypePerInt4; ++ j)
         out_dtypes[j] = static_cast<dtype_t>(values[j]);
     recv_int4[token_idx * hidden_int4 + i] = out_int4;
 }

enum class WarpRole {
        kNVLSender, //从NVLINK上发送
        kNVLAndRDMAForwarder,//NVLink到RDMA转发
        kRDMAReceiver, //RDMA接收
        kCoordinator //协调器
    };

const auto rdma_rank = rank / NUM_MAX_NVL_PEERS, nvl_rank = rank % NUM_MAX_NVL_PEERS;
    auto role_meta = [=]() -> std::pair<WarpRole, int> {
        auto warp_id = thread_id / 32;
        if (not is_rdma_receiver_sm) {
            if (warp_id < NUM_MAX_NVL_PEERS) {
                auto shuffled_warp_id = warp_id;
                shuffled_warp_id = (shuffled_warp_id + channel_id) % NUM_MAX_NVL_PEERS;
                return {WarpRole::kNVLSender, shuffled_warp_id};
            } elseif (warp_id < NUM_MAX_NVL_PEERS + kNumForwarders) {
                auto shuffled_warp_id = warp_id - NUM_MAX_NVL_PEERS;
                shuffled_warp_id = (shuffled_warp_id + channel_id) % kNumForwarders;
                return {WarpRole::kNVLAndRDMAForwarder, shuffled_warp_id};
            } else {
                return {WarpRole::kCoordinator, 0};
            }
        } else {
            if (warp_id < NUM_MAX_NVL_PEERS + kNumForwarders) {
                return {WarpRole::kRDMAReceiver, warp_id};
            } else {
                return {WarpRole::kCoordinator, 0};
            }
        }
    }();

struct LowLatencyBuffer {
    int num_clean_int = 0;

    void* dispatch_rdma_send_buffer = nullptr;
    void* dispatch_rdma_recv_data_buffer = nullptr;
    int* dispatch_rdma_recv_count_buffer = nullptr;
    int* dispatch_rdma_atomic_token_counter = nullptr;

    void* combine_rdma_send_buffer = nullptr;
    void* combine_rdma_recv_data_buffer = nullptr;
    int* combine_rdma_recv_flag_buffer = nullptr;

    std::pair<int*, int> clean_meta() {
        EP_HOST_ASSERT(dispatch_rdma_recv_count_buffer == combine_rdma_recv_flag_buffer);
        return {dispatch_rdma_recv_count_buffer, num_clean_int};
    }
};

struct LowLatencyLayout {
    size_t total_bytes = 0;
    LowLatencyBuffer buffers[2];

    template <typenameout_ptr_t = void*, typenamecount_ptr_t = uint8_t*, typenamein_ptr_t = void*>
    out_ptr_t advance(constin_ptr_t& ptr, size_t count) {
        returnreinterpret_cast<out_ptr_t>(reinterpret_cast<count_ptr_t>(ptr) + count);
    }

    LowLatencyLayout(void* rdma_buffer, int num_max_dispatch_tokens_per_rank, int hidden, int num_ranks, int num_experts) {
        constint num_scales = hidden / 128;
        constint num_local_experts = num_experts / num_ranks;

        // Dispatch and combine layout:
        //  - 2 symmetric odd/even send buffer
        //  - 2 symmetric odd/even receive buffers
        //  - 2 symmetric odd/even signaling buffers

// Message sizes
        EP_HOST_ASSERT(num_scales * sizeof(float) <= hidden);
        size_t num_bytes_per_dispatch_msg = hidden + num_scales * sizeof(float) + sizeof(int4);
        size_t num_bytes_per_combine_msg = sizeof(int4) + hidden * sizeof(nv_bfloat16);

// Send buffer
        size_t dispatch_send_buffer_bytes = num_max_dispatch_tokens_per_rank * num_bytes_per_dispatch_msg;
        size_t combine_send_buffer_bytes = num_experts * num_max_dispatch_tokens_per_rank * num_bytes_per_combine_msg;
        size_t send_buffer_bytes = std::max(dispatch_send_buffer_bytes, combine_send_buffer_bytes);
        EP_HOST_ASSERT(send_buffer_bytes % sizeof(int4) == 0);
        total_bytes += send_buffer_bytes * 2;

        // Symmetric receive buffers
        // TODO: optimize memory usages
        size_t dispatch_recv_data_buffer_bytes = num_experts * num_max_dispatch_tokens_per_rank * num_bytes_per_dispatch_msg;
        size_t combine_recv_buffer_bytes = num_experts * num_max_dispatch_tokens_per_rank * num_bytes_per_combine_msg;
        size_t recv_buffer_bytes = std::max(dispatch_recv_data_buffer_bytes, combine_recv_buffer_bytes);
        EP_HOST_ASSERT(recv_buffer_bytes % sizeof(int4) == 0);
        total_bytes += recv_buffer_bytes * 2;

        // Symmetric signaling buffers
        size_t dispatch_recv_count_buffer_bytes = num_experts * sizeof(int);
        size_t dispatch_recv_atomic_token_counter_bytes = num_local_experts * sizeof(int);
        size_t combine_recv_flag_buffer_bytes = dispatch_recv_count_buffer_bytes;
        size_t signaling_buffer_bytes = std::max(dispatch_recv_count_buffer_bytes + dispatch_recv_atomic_token_counter_bytes,
                                                 combine_recv_flag_buffer_bytes);
        total_bytes += signaling_buffer_bytes * 2;

def low_latency_dispatch(hidden_states: torch.Tensor, topk_idx: torch.Tensor, num_max_dispatch_tokens_per_rank: int, num_experts: int):
    global _buffer

    # Do MoE dispatch, compatible with CUDA graph (but you may restore some buffer status once you replay)
    recv_hidden_states, recv_expert_count, handle, event, hook = \
        _buffer.low_latency_dispatch(hidden_states, topk_idx, num_max_dispatch_tokens_per_rank, num_experts,
                                     async_finish=False, return_recv_hook=True)

    # NOTES: the actual tensor will not be received only if you call `hook()`,
    # it is useful for double-batch overlapping, but **without any SM occupation**
    # If you don't want to overlap, please set `return_recv_hook=False`
    # Later, you can use our GEMM library to do the computation with this specific format
    return recv_hidden_states, recv_expert_count, handle, event, hook

// Allocate packed tensors
    auto packed_recv_x = torch::empty({num_local_experts, num_ranks * num_max_dispatch_tokens_per_rank, hidden}, x.options().dtype(torch::kFloat8_e4m3fn));
    auto packed_recv_src_info = torch::empty({num_local_experts, num_ranks * num_max_dispatch_tokens_per_rank}, torch::dtype(torch::kInt32).device(torch::kCUDA));
    auto packed_recv_layout_range = torch::empty({num_local_experts, num_ranks}, torch::dtype(torch::kInt64).device(torch::kCUDA));
    auto packed_recv_count = torch::from_blob(buffer.dispatch_rdma_atomic_token_counter,
                                              {num_local_experts}, torch::dtype(torch::kInt32).device(torch::kCUDA));

// Allocate column-majored scales
    EP_HOST_ASSERT((num_ranks * num_max_dispatch_tokens_per_rank) % 4 == 0 and "TMA requires the number of tokens to be multiple of 4");
    auto packed_recv_x_scales = torch::empty({num_local_experts, num_scales, num_ranks * num_max_dispatch_tokens_per_rank}, torch::dtype(torch::kFloat32).device(torch::kCUDA));
    packed_recv_x_scales = torch::transpose(packed_recv_x_scales, 1, 2);

// FP8 cast
#pragma unroll
for (int i = thread_id; i < hidden_bf16_int4; i += num_threads) {
    // Read and calculate local amax
    auto int4_value = __ldg(x_int4 + i);
    auto bf16_values = reinterpret_cast<nv_bfloat16*>(&int4_value);
    float fp32_values[kNumElemsPerRead];
    float amax = kFP8Margin, scale, scale_inv;
    #pragma unroll
    for (int j = 0; j < kNumElemsPerRead; ++ j) {
        fp32_values[j] = static_cast<float>(bf16_values[j]);
        amax = fmaxf(amax, fabsf(fp32_values[j]));
    }

    // Reduce amax and scale
    EP_STATIC_ASSERT(kNumElemsPerRead * 32 / kNumPerChannels == 2, "Invalid vectorization");
    amax = half_warp_reduce_max(amax), scale = kFP8Amax / amax, scale_inv = amax * kFP8AmaxInv;
    if (lane_id == 0or lane_id == 16)
        rdma_x_scales[i * kNumElemsPerRead / 128] = scale_inv;

    // Cast into send buffer
    int2 int2_value;
    auto fp8x2_values = reinterpret_cast<__nv_fp8x2_storage_t*>(&int2_value);
    #pragma unroll
    for (int j = 0; j < kNumElemsPerRead; j += 2) {
        float2 fp32x2 = {fp32_values[j] * scale, fp32_values[j + 1] * scale};
        fp8x2_values[j / 2] = __nv_cvt_float2_to_fp8x2(fp32x2, __NV_SATFINITE, __NV_E4M3);
    }
    rdma_x_int2[i] = int2_value;
}

// Issue IBGDA sends
if (dst_expert_idx >= 0) {
    int slot_idx = lane_id == 0 ? atomicAdd(atomic_counter_per_expert + dst_expert_idx, 1) : 0;
    slot_idx = __shfl_sync(0xffffffff, slot_idx, 0);
    constauto dst_rank = dst_expert_idx / num_local_experts;
    constauto dst_expert_local_idx = dst_expert_idx % num_local_experts;
    constauto src_ptr = reinterpret_cast<uint64_t>(rdma_x_int2);
    constauto dst_ptr = reinterpret_cast<uint64_t>(rdma_recv_x) +
                         dst_expert_local_idx * num_ranks * num_max_dispatch_tokens_per_rank * num_bytes_per_msg +
                         rank * num_max_dispatch_tokens_per_rank * num_bytes_per_msg +
                         slot_idx * num_bytes_per_msg;
    if (dst_rank != rank) {
        nvshmemi_ibgda_put_nbi_warp(dst_ptr, src_ptr, num_bytes_per_msg, dst_rank, dst_expert_local_idx, lane_id, slot_idx);
    } else {
        // NOTES: only 2 load iterations for 7K hidden with 8 unrolls
        constauto* src_int4_ptr = reinterpret_cast<const int4*>(src_ptr);
        constauto* dst_int4_ptr = reinterpret_cast<int4*>(dst_ptr);
        UNROLLED_WARP_COPY(8, lane_id, num_int4_per_msg, dst_int4_ptr, src_int4_ptr, ld_nc_global, st_na_global);
    }
    // Increase counter after finishing
    __syncwarp();
    lane_id == 0 ? atomic_add_release_global(atomic_finish_counter_per_expert + dst_expert_idx, 1) : 0;
}

def low_latency_combine(hidden_states: torch.Tensor,
                        topk_idx: torch.Tensor, topk_weights: torch.Tensor, handle: Tuple):
    global _buffer

    # Do MoE combine, compatible with CUDA graph (but you may restore some buffer status once you replay)
    combined_hidden_states, event_overlap, hook = \
        _buffer.low_latency_combine(hidden_states, topk_idx, topk_weights, handle,
                                    async_finish=False, return_recv_hook=True)

    # NOTES: the same behavior as described in the dispatch kernel
    return combined_hidden_states, event_overlap, hook

// FP8 cast and issue IBGDA sends
    if (responsible_expert_idx < num_experts) {
        constauto dst_rank = responsible_expert_idx / num_local_experts;
        constauto local_expert_idx = responsible_expert_idx % num_local_experts;
        constauto global_expert_idx = rank * num_local_experts + local_expert_idx;
        constauto layout = __ldg(layout_range + local_expert_idx * num_ranks + dst_rank);
        constauto local_x = reinterpret_cast<const int4*>(x) +
                local_expert_idx * num_ranks * num_max_dispatch_tokens_per_rank * hidden_bf16_int4;
        constauto local_src_info = src_info + local_expert_idx * num_ranks * num_max_dispatch_tokens_per_rank;
        constauto rdma_send_x_vec = reinterpret_cast<uint8_t*>(rdma_send_x) +
                local_expert_idx * num_ranks * num_max_dispatch_tokens_per_rank * num_bytes_per_slot;

        // Unpack layout
        int offset, num_tokens_to_send;
        unpack2(layout, num_tokens_to_send, offset);

        // Issue IBGDA send
        for (int token_idx = offset + sub_warp_id; token_idx < offset + num_tokens_to_send; token_idx += kNumWarpsPerGroup) {
            constauto x_int4 = local_x + token_idx * hidden_bf16_int4;
            constauto rdma_send_type_row = reinterpret_cast<int*>(rdma_send_x_vec + token_idx * num_bytes_per_slot);
            constauto rdma_send_x_vec_row = reinterpret_cast<uint8_t*>(rdma_send_type_row + 4);

            // Copy directly to local rank, or copy to buffer and issue RDMA
            auto src_idx = __ldg(local_src_info + token_idx);
            constauto buf_ptr = reinterpret_cast<int64_t>(rdma_send_x_vec_row);
            constauto dst_ptr = reinterpret_cast<uint64_t>(rdma_recv_x) + (global_expert_idx * num_max_dispatch_tokens_per_rank + src_idx) * num_bytes_per_slot + sizeof(int4);
            if (dst_rank == rank) {
                constauto dst_int4_ptr = reinterpret_cast<int4*>(dst_ptr);
                UNROLLED_WARP_COPY(7, lane_id, hidden_bf16_int4, dst_int4_ptr, x_int4, ld_nc_global, st_na_global);
            } else {
                constauto buf_int4_ptr = reinterpret_cast<int4*>(buf_ptr);
                UNROLLED_WARP_COPY(7, lane_id, hidden_bf16_int4, buf_int4_ptr, x_int4, ld_nc_global, st_na_global);
                nvshmemi_ibgda_put_nbi_warp(dst_ptr, buf_ptr, hidden * sizeof(nv_bfloat16), dst_rank, local_expert_idx, lane_id, token_idx - offset);
            }
        }

        // Put finishing flag
        EP_STATIC_ASSERT(kNumWarpsPerGroup > 1, "Requires more than one warp per group");
        asm volatile("bar.sync %0, %1;" :: "r"(warp_group_id + 1), "r"(kNumWarpsPerGroup * 32));
        if (sub_warp_id == 1and lane_id == 0) {
            while (ld_acquire_global(atomic_clean_flag) == 0);
            if (dst_rank != rank) {
                nvshmemi_ibgda_rma_p(rdma_recv_flag + global_expert_idx, 1, dst_rank, local_expert_idx, 0);
            } else {
                st_na_release(rdma_recv_flag + global_expert_idx, 1);
            }
            atomic_add_release_global(atomic_clean_flag, -1);
        }
        __syncwarp();
    }

#ifndef DISABLE_AGGRESSIVE_PTX_INSTRS
#define LD_NC_FUNC "ld.global.nc.L1::no_allocate.L2::256B"
#else
#define LD_NC_FUNC "ld.volatile.global"
#endif

// `ld.global.nc.L1::no_allocate` will be translated into `LDG.E.NA.[width].CONSTANT` in SASS,
// which does not have cache allocation, and `CONSTANT` memory does not have coherence control,
// so we have to control them by queue semantics
template <typenamedtype_t>
__device__  __forceinline__ dtype_t ld_nc_global(const dtype_t *ptr) {
    auto ret = ld_nc_global(reinterpret_cast<consttypename VecInt<sizeof(dtype_t)>::vec_t*>(ptr));
    return *reinterpret_cast<dtype_t*>(&ret);
}