Add an operation to perform splitting for context parallelism. This operation split the input_ids into cp_size chunks, and return the cp_rank-th chunk. When the seqlen % cp_size != 0, the chunk sizes of each rank would be [seqlen // cp_size, seqlen // cp_size, ..., seqlen - (se
(
input_ids: Tensor,
host_request_types: Tensor,
host_context_lengths: Tensor, # for pad-free input mode
cp_size: int = 1,
cp_rank: int = 0,
)
| 7586 | |
| 7587 | |
| 7588 | def cp_split_plugin( |
| 7589 | input_ids: Tensor, |
| 7590 | host_request_types: Tensor, |
| 7591 | host_context_lengths: Tensor, # for pad-free input mode |
| 7592 | cp_size: int = 1, |
| 7593 | cp_rank: int = 0, |
| 7594 | ) -> Tensor: |
| 7595 | ''' |
| 7596 | Add an operation to perform splitting for context parallelism. |
| 7597 | |
| 7598 | This operation split the input_ids into cp_size chunks, and return the cp_rank-th |
| 7599 | chunk. |
| 7600 | When the seqlen % cp_size != 0, the chunk sizes of each rank would be |
| 7601 | [seqlen // cp_size, seqlen // cp_size, ..., seqlen - (seqlen // cp_size) * cp_size] |
| 7602 | |
| 7603 | It inserts a IPluginV3Layer. |
| 7604 | |
| 7605 | Parameters: |
| 7606 | input : Tensor |
| 7607 | The input tensor contains the indices to split. |
| 7608 | |
| 7609 | host_request_types: Tensor = None (On CPU) |
| 7610 | The tensor on the host that indicates if a request is in context or |
| 7611 | generation phase. Its shape is [batch_size]. See Inflight Batching |
| 7612 | in docs/gpt_attention.md, |
| 7613 | |
| 7614 | host_context_lengths: Tensor = None (On CPU) |
| 7615 | A host tensor that contains the lengths of the different inputs |
| 7616 | |
| 7617 | Returns: |
| 7618 | The output split tensor. |
| 7619 | The length of the output split tensor. |
| 7620 | The index for rebuilding the sequence |
| 7621 | ''' |
| 7622 | plg_creator = trt.get_plugin_registry().get_creator( |
| 7623 | 'CpSplit', '1', TRT_LLM_PLUGIN_NAMESPACE) |
| 7624 | assert plg_creator is not None |
| 7625 | |
| 7626 | cp_size = trt.PluginField("cp_size", np.array([int(cp_size)], np.int32), |
| 7627 | trt.PluginFieldType.INT32) |
| 7628 | cp_rank = trt.PluginField("cp_rank", np.array([int(cp_rank)], np.int32), |
| 7629 | trt.PluginFieldType.INT32) |
| 7630 | |
| 7631 | pfc = trt.PluginFieldCollection([cp_size, cp_rank]) |
| 7632 | cp_split_plug = plg_creator.create_plugin("cp_split", pfc, |
| 7633 | trt.TensorRTPhase.BUILD) |
| 7634 | plug_inputs = [ |
| 7635 | input_ids.trt_tensor, host_request_types.trt_tensor, |
| 7636 | host_context_lengths.trt_tensor |
| 7637 | ] |
| 7638 | |
| 7639 | layer = default_trtnet().add_plugin_v3(plug_inputs, [], cp_split_plug) |
| 7640 | _add_plugin_info(layer, plg_creator, "cp_split", pfc) |
| 7641 | return _create_tensor(layer.get_output(0), |
| 7642 | layer), _create_tensor(layer.get_output(2), layer) |
no test coverage detected