Function recv

tensorrt_llm/functional.py:4298–4338

Add an operation that performs a recv to a rank from another. The recv operation receives a tensor from on a rank from another. If a rank 'i' receives a tensor from a rank 'j', the rank 'j' must have a corresponding 'send' operation to rank 'j'. See 'send'. That operation is i

(tensor: Tensor, src: int)

def recv(tensor: Tensor, src: int) -> Tensor:
    '''
    Add an operation that performs a recv to a rank from another.

    The recv operation receives a tensor from on a rank from another. If a rank 'i'
    receives a tensor from a rank 'j', the rank 'j' must have a corresponding 'send'
    operation to rank 'j'. See 'send'.

    That operation is implemented using a plugin that wraps the NCCL recv
    point-to-point operation. See
    https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/api/p2p.html#ncclrecv
    for details.

    Parameters:
        tensor : Tensor
            The input tensor.

        src : int
            The rank that sends the tensor to.

    Returns:
        The tensor produced by that layer.
    '''
    recv_plg_creator = trt.get_plugin_registry().get_plugin_creator(
        'Recv', '1', TRT_LLM_PLUGIN_NAMESPACE)
    assert recv_plg_creator is not None

    src = trt.PluginField("src_rank", np.array(src, dtype=np.int32),
                          trt.PluginFieldType.INT32)
    p_dtype = default_net().plugin_config.nccl_plugin
    pf_type = trt.PluginField(
        "type_id", np.array([int(str_dtype_to_trt(p_dtype))], np.int32),
        trt.PluginFieldType.INT32)

    pfc = trt.PluginFieldCollection([src, pf_type])
    recv_plug = recv_plg_creator.create_plugin("recv", pfc)
    plug_inputs = [tensor.cast(p_dtype).trt_tensor]

    layer = default_trtnet().add_plugin_v2(plug_inputs, recv_plug)
    _add_plugin_info(layer, recv_plg_creator, "recv", pfc)
    return _create_tensor(layer.get_output(0), layer).cast(tensor.dtype)


def gemm_allreduce(a: Tensor,