flashinfer.comm.vllm_all_reduce

flashinfer.comm.vllm_all_reduce(fa: int, inp: Tensor, out: Tensor, reg_buffer: int, reg_buffer_sz_bytes: int, num_ctas: int) → None

执行一个就地 all reduce。

参数:

• fa – 自定义 all reduce 的句柄。

• inp – 要 all reduce 的输入张量。

• out – 要 all reduce 的输出张量。

• reg_buffer – 要 all reduce 的寄存器缓冲区。

• reg_buffer_sz_bytes – 寄存器缓冲区的大小（字节）。

• num_ctas – all reduce 要使用的 CTA 数量。

• bounds (CTA 上限) –

  36. 通常，即使使用少量 SM，我们也可以饱和带宽。