【cuda】get_low_latency_rdma_size_hint

好的我们来详细分析get_low_latency_rdma_size_hint这个函数是如何计算所需 RDMA 缓冲区大小的。首先我们需要理解在低延迟模式 (Low-latency Mode)下RDMA 缓冲区是如何被使用的。低延迟模式下 RDMA 缓冲区的布局和用途在低延迟模式中通信模型是完全端到端的。每个 GPU 都需要有能力向其他任何一个 GPU 发送数据并从其他任何一个 GPU 接收数据。为了避免复杂的实时内存管理和同步deep_ep采用了一种静态预分配的策略。RDMA 缓冲区被划分成两个主要部分实现了双缓冲 (Double Buffering)或乒乓缓冲 (Ping-Pong Buffering)我们称之为Buffer 0和Buffer 1。这允许在一个缓冲区上进行计算或数据处理同时在另一个缓冲区上进行网络通信。每个独立的缓冲区比如Buffer 0又被进一步划分为几个功能区分发-发送区 (dispatch_send_buffer):用途: 当前 GPU (rank_me) 准备要发送给所有其他 GPU的数据。布局: 它会为要发送给rank_0,rank_1, …,rank_N-1的每个目标 GPU 都预留一小块空间。大小: 需要能容纳num_max_dispatch_tokens_per_rank个 token。分发-接收区 (dispatch_recv_buffer):用途: 存放从所有其他 GPU接收到的数据。布局: 它会为来自rank_0,rank_1, …,rank_N-1的每个源 GPU 都预留一小块空间。大小: 每个小块需要能容纳num_max_dispatch_tokens_per_rank个 token。计数区: 旁边还有一个区域用于存放每个源 GPU 实际发送了多少个 token 的计数。合并-发送区 (combine_send_buffer):用途: 当前 GPU 在专家计算完成后准备要发送回原始 GPU 的数据。布局: 与dispatch_send_buffer类似为每个目标 GPU 预留空间。合并-接收区 (combine_recv_buffer):用途: 存放从其他 GPU reduce 回来的最终结果。布局: 与dispatch_recv_buffer类似为每个源 GPU 预留空间。标志区: 旁边有一个区域用于原子地标记数据是否已到达。get_low_latency_rdma_size_hint的任务就是计算单个缓冲区比如Buffer 0所包含的这些功能区的总大小。最终的总 RDMA 缓冲区大小会是这个值的两倍因为有Buffer 0和Buffer 1。C 源码分析让我们来看deep_ep.cpp中对应的 C 实现。计算逻辑封装在LowLatencyLayout这个结构体中。// in csrc/kernels/utils.cuh (or similar file)structLowLatencyLayout{// ...LowLatencyLayout(void*base_ptr,intnum_max_dispatch_tokens_per_rank,inthidden,intnum_ranks,intnum_experts){// ... (一些断言) ...// 计算每个 消息 (一个 token 元数据) 的大小intnum_bytes_per_dispatch_msghidden*sizeof(bfloat16)sizeof(int);// 数据 源信息 (src_info)num_bytes_per_dispatch_msg_alignedalign_to(num_bytes_per_dispatch_msg,16);intnum_bytes_per_combine_msghidden*sizeof(bfloat16);// 只需要数据num_bytes_per_combine_msg_alignedalign_to(num_bytes_per_combine_msg,16);// --- 计算 Buffer 0 的布局 ---// 1. Dispatch Recv Data Bufferint64_tcurrent_offset0;int64_tnum_bytes_dispatch_recv_data_buffer(int64_t)num_max_dispatch_tokens_per_rank*num_ranks*num_bytes_per_dispatch_msg_aligned;current_offsetnum_bytes_dispatch_recv_data_buffer;// 2. Dispatch Recv Count Bufferint64_tnum_bytes_dispatch_recv_count_buffer(int64_t)num_ranks*sizeof(int);current_offsetnum_bytes_dispatch_recv_count_buffer;// 3. Dispatch Send Buffer// (注意发送缓冲区在逻辑上存在但在物理上可以和接收缓冲区复用// 因为发送的数据会被直接 RDMA put 到目标GPU的接收区// 但为了简化计算我们假设它需要独立空间)int64_tnum_bytes_dispatch_send_buffer(int64_t)num_max_dispatch_tokens_per_rank*num_ranks*num_bytes_per_dispatch_msg_aligned;current_offsetnum_bytes_dispatch_send_buffer;// 4. Combine Recv Data Bufferint64_tnum_bytes_combine_recv_data_buffer(int64_t)num_max_dispatch_tokens_per_rank*num_ranks*num_bytes_per_combine_msg_aligned;current_offsetnum_bytes_combine_recv_data_buffer;// 5. Combine Recv Flag Bufferint64_tnum_bytes_combine_recv_flag_buffer(int64_t)num_ranks*num_max_dispatch_tokens_per_rank*sizeof(int);current_offsetnum_bytes_combine_recv_flag_buffer;// 6. Combine Send Bufferint64_tnum_bytes_combine_send_buffer(int64_t)num_max_dispatch_tokens_per_rank*num_ranks*num_bytes_per_combine_msg_aligned;current_offsetnum_bytes_combine_send_buffer;// ... (对 current_offset 进行对齐) ...// 单个乒乓缓冲区的总大小int64_tsingle_buffer_bytesalign_to(current_offset,4096);// 乘以2得到双缓冲的总大小total_bytessingle_buffer_bytes*2;}};// python 绑定的函数int64_tget_low_latency_rdma_size_hint(...){LowLatencyLayoutlayout(nullptr,...);// 传入nullptr只为了计算大小returnlayout.total_bytes;}注意: 上述 C 代码是一个为了解释而简化的版本实际的LowLatencyLayout实现可能会更复杂例如通过内存复用等技巧来优化空间。但其计算总量的核心逻辑是一致的。详细举例计算让我们使用test_low_latency.py中的默认参数来计算num_tokens(num_max_dispatch_tokens_per_rank) 128hidden 7168num_ranks 8num_experts 288bfloat16大小 2 字节1. 计算单个消息的大小Dispatch 消息:每个 token 的数据:hidden * sizeof(bfloat16) 7168 * 2 14336 字节每个 token 的元数据 (src_info):sizeof(int) 4 字节num_bytes_per_dispatch_msg 14336 4 14340 字节对齐到 16 字节:align_to(14340, 16) 14352 字节Combine 消息:每个 token 的数据:hidden * sizeof(bfloat16) 7168 * 2 14336 字节对齐到 16 字节:align_to(14336, 16) 14336 字节 (本身已对齐)2. 计算单个乒乓缓冲区 (Buffer 0) 的大小Dispatch Recv Data Buffer: 用于存放从所有num_ranks个 GPU 收到的、最多num_tokens个 token。size num_tokens * num_ranks * num_bytes_per_dispatch_msg_alignedsize 128 * 8 * 14352 14,696,448 字节 (约 14.0 MB)Dispatch Recv Count Buffer: 用于存放从每个num_ranksGPU 到底收到了多少 token。size num_ranks * sizeof(int)size 8 * 4 32 字节Combine Recv Data Buffer: 用于存放从所有num_ranks个 GPU reduce 回来的结果。size num_tokens * num_ranks * num_bytes_per_combine_msg_alignedsize 128 * 8 * 14336 14,680,064 字节 (约 14.0 MB)Combine Recv Flag Buffer: 需要为可能收到的每个 token 准备一个标志位。size num_ranks * num_tokens * sizeof(int)size 8 * 128 * 4 4,096 字节Send Buffers: 在最坏的情况下发送缓冲区和接收缓冲区大小相同。所以我们简单地将接收区的大小乘以2。dispatch_total (dispatch_recv_data dispatch_recv_count) * 2combine_total (combine_recv_data combine_recv_flag) * 2总和 (简化估算):total_single_buffer_approx dispatch_recv_data * 2 combine_recv_data * 2total_single_buffer_approx ≈ 14.0 MB * 2 14.0 MB * 2 56 MB这只是一个粗略估算实际计算会更精确并包含所有小块。3. 计算最终总大小对齐: 将单个缓冲区的总大小向上对齐到一个较大的边界比如 4KB 页面大小。双缓冲: 将对齐后的单个缓冲区大小乘以 2。num_rdma_bytes align_to(total_single_buffer_bytes, 4096) * 2num_rdma_bytes ≈ 56 MB * 2 112 MB所以最终get_low_latency_rdma_size_hint返回的值大约在 112 MB 左右。这个值确保了在最坏的情况下每个 GPU 都需要向其他所有 GPU 发送/接收最大数量的 tokenRDMA 缓冲区也绝对够用不会发生内存越界。这个**“hint”**之所以是“提示”是因为它基于最坏情况的静态分析实际使用中可能并不总是需要这么多空间但预分配这么多可以极大地简化运行时逻辑避免动态分配和复杂的同步从而实现低延迟。