flashinfer.comm.moe_a2a_wrap_payload_tensor_in_workspace

flashinfer.comm.moe_a2a_wrap_payload_tensor_in_workspace(workspace: Tensor, leading_shape: list[int], slice_start: int, slice_end: int, dtype: dtype) → Tensor

将工作区中的偏移量包装成一个张量。

参数:

• workspace – [ep_size, size_per_rank] 或 [size_per_rank] 工作区张量

• leading_shape – 用于包装张量的前导形状

• slice_start – 工作区中切片的起始位置

• slice_end – 工作区中切片的结束位置

• dtype – 输出张量的数据类型

返回值:

[leading_shape, *] 工作区支持的张量

返回值类型:

张量