Add an operation that performs a send from a rank to another. The send operation sends a tensor from one rank to another. If a rank 'i' sends a tensor to a rank 'j', the rank 'j' must have a corresponding 'recv' operation from rank 'i'. See 'recv'. That operation is implemente
(tensor: Tensor, tgt: int)
| 4252 | |
| 4253 | |
| 4254 | def send(tensor: Tensor, tgt: int) -> Tensor: |
| 4255 | ''' |
| 4256 | Add an operation that performs a send from a rank to another. |
| 4257 | |
| 4258 | The send operation sends a tensor from one rank to another. If a rank 'i' |
| 4259 | sends a tensor to a rank 'j', the rank 'j' must have a corresponding 'recv' |
| 4260 | operation from rank 'i'. See 'recv'. |
| 4261 | |
| 4262 | That operation is implemented using a plugin that wraps the NCCL send |
| 4263 | point-to-point operation. See |
| 4264 | https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/api/p2p.html#ncclsend |
| 4265 | for details. |
| 4266 | |
| 4267 | Parameters: |
| 4268 | tensor : Tensor |
| 4269 | The input tensor. |
| 4270 | |
| 4271 | tgt : int |
| 4272 | The rank that receives the tensor. |
| 4273 | |
| 4274 | Returns: |
| 4275 | The tensor produced by that layer. |
| 4276 | ''' |
| 4277 | send_plg_creator = trt.get_plugin_registry().get_plugin_creator( |
| 4278 | 'Send', '1', TRT_LLM_PLUGIN_NAMESPACE) |
| 4279 | assert send_plg_creator is not None |
| 4280 | |
| 4281 | tgt = trt.PluginField("tgt_rank", np.array(tgt, dtype=np.int32), |
| 4282 | trt.PluginFieldType.INT32) |
| 4283 | |
| 4284 | p_dtype = default_net().plugin_config.nccl_plugin |
| 4285 | pf_type = trt.PluginField( |
| 4286 | "type_id", np.array([int(str_dtype_to_trt(p_dtype))], np.int32), |
| 4287 | trt.PluginFieldType.INT32) |
| 4288 | |
| 4289 | pfc = trt.PluginFieldCollection([tgt, pf_type]) |
| 4290 | send_plug = send_plg_creator.create_plugin("send", pfc) |
| 4291 | plug_inputs = [tensor.cast(p_dtype).trt_tensor] |
| 4292 | |
| 4293 | layer = default_trtnet().add_plugin_v2(plug_inputs, send_plug) |
| 4294 | _add_plugin_info(layer, send_plg_creator, "send", pfc) |
| 4295 | return _create_tensor(layer.get_output(0), layer).cast(tensor.dtype) |
| 4296 | |
| 4297 | |
| 4298 | def recv(tensor: Tensor, src: int) -> Tensor: |