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¶

Wrap an offset in the workspace into a tensor.

Parameters:

workspace – [ep_size, size_per_rank] or [size_per_rank] workspace tensor
leading_shape – The leading shape to wrap the tensor with
slice_start – The start of the slice in the workspace
slice_end – The end of the slice in the workspace
dtype – Data type for the output tensor

Returns:

[leading_shape, *] workspace-backed tensor

Return type:

tensor