【vllm】deepep

auto recv_x torch::empty({num_recv_tokens, hidden}, x.options());这行代码的作用是为当前 GPU 分配一块显存缓冲区用于存放即将从其他所有 GPU 接收到的 token 数据。这是dispatch操作中一个至关重要的步骤我们可以把它理解为“为客人准备好房间”。详细解释让我们把dispatch操作想象成一个“派对”每个 GPU 都是派对的组织者同时也扮演着客人的角色。x张量: 这是当前 GPU 自己拥有的、准备分发出去的原始 token 数据。可以看作是“我要带去派对的礼物”。dispatch的目标: 将x中需要发送给其他 GPU 的 token通过网络NVLink/RDMA发送出去。同时接收其他 GPU 发送给自己的 token。num_recv_tokens: 在这行代码被执行之前deep_ep已经通过notify_dispatch阶段完成了一轮“尺寸元数据”的交换。每个 GPU 都向全网广播了它将要接收多少个 token。通过对这些信息进行all_reduce求和每个 GPU 都精确地计算出了自己总共将要从所有其他 GPU 接收到的 token 数量这个总数就是num_recv_tokens。可以理解为派对开始前每个组织者都已经统计好了“总共会有多少位客人来我这里”。torch::empty({num_recv_tokens, hidden}, x.options()): 这就是“准备房间”的核心操作。torch::empty(...): 这是 PyTorch C API 中用于分配一块未初始化的 GPU 内存的函数。使用empty而不是zeros是为了性能因为我们知道这块内存马上就会被从网络接收到的数据完全覆盖没有必要预先清零。{num_recv_tokens, hidden}: 这是新分配的缓冲区的形状 (Shape)。num_recv_tokens: 缓冲区的行数正好等于我们要接收的 token 总数。不多不少精确匹配。hidden: 缓冲区的列数与原始 token 的隐藏层维度一致。x.options(): 这是一个方便的写法它告诉torch::empty新创建的张量recv_x应该具有与输入张量x相同的属性包括dtype: 数据类型例如torch::kBFloat16。device: 所在的设备例如torch::kCUDA。layout: 内存布局例如torch::kStrided(表示是稠密的)。最终效果执行完这行代码后在当前 GPU 的显存中一块名为recv_x的、大小为[num_recv_tokens, hidden]的连续内存区域被成功分配。这块内存已经准备就绪可以作为后续dispatchCUDA Kernel 的写入目标 (Destination)。当dispatchKernel 运行时它会从网络NVLink/RDMA 缓冲区中读取数据并直接将接收到的 token 一个个地填充到recv_x这块预留好的“客房”里。为什么必须先分配CUDA Kernel 和网络硬件如 RDMA 网卡不能凭空写入内存它们需要一个明确的、合法的、已经分配好的目标内存地址。这行代码就是提供了这个必需的目标地址。如果事先不知道要接收多少数据 (num_recv_tokens)就无法精确地分配缓冲区这会导致要么内存浪费分配过大要么内存溢出分配过小。因此notify_dispatch阶段的尺寸交换是dispatch操作能够高效、安全执行的前提。