Class InputMetaData

colossalai/inference/config.py:53–147

The input info for a single step Args: block_tables (torch.Tensor, optional): Sequences' BlockTables Defaults to None. sequence_lengths (torch.Tensor): A tensor containing sequence lengths. fd_inter_tensor (torch.Tensor, optional): A tensor representing intermediate data for flash d

@dataclass
class InputMetaData(RPC_PARAM):
    """The input info for a single step

    Args:
    block_tables (torch.Tensor, optional): Sequences' BlockTables Defaults to None.
    sequence_lengths (torch.Tensor): A tensor containing sequence lengths.
    fd_inter_tensor (torch.Tensor, optional): A tensor representing intermediate data for flash decoding. Defaults to None.
    batch_size (int, optional): The current batch size. Defaults to 64.
    is_prompts (bool, optional): Indicates whether prefill or decoding. Defaults to False(decoding).
    use_cuda_kernel(bool): Whether to use cuda kernel, faster but lose some precision occasionally
    use_cuda_graph (bool, optional): Indicates whether to use the CUDA graph. Defaults to False.
    kv_seq_len (int, optional): Key-value sequence length. Defaults to 512.
    head_dim (int, optional): Head dimension. Defaults to 32.
    high_precision(bool, optional): Whether to use float32 for underlying calculations of float16 data to achieve higher precision, Defaults to False.
    dtype (torch.dtype, optional): The computation type of tensor, Defaults to torch.float32.
    use_spec_dec (bool): Indicate whether to use speculative decoding.
    num_tokens_to_verify (int): The number of tokens to verify in speculative decoding. Only valid when `use_spec_dec` is set to True.
    batch_token_ids (List[List[int]], optional): input_token_ids + output_token_ids of current batch. Only used for `repetition_penalty`, `no_repeat_ngram_size` in sampler process.
    """

    block_tables: torch.Tensor = None
    sequence_lengths: torch.Tensor = None
    fd_inter_tensor: FDIntermTensors = None
    batch_size: int = 64  # current_batch_size
    is_prompts: bool = False
    use_cuda_kernel: bool = False
    use_cuda_graph: bool = False
    kv_seq_len: int = 512
    head_dim: int = 32
    high_precision: bool = False
    dtype: torch.dtype = torch.float32
    use_spec_dec: bool = False
    num_tokens_to_verify: int = 0
    batch_token_ids: Optional[List[List[int]]] = (
        None  # for `repetition_penalty`, `no_repeat_ngram_size` in sampler process
    )

    def to_rpc_param(self) -> Dict[str, any]:
        return {
            "block_tables": self.block_tables.tolist(),
            "sequence_lengths": self.sequence_lengths.tolist(),
            "batch_size": self.batch_size,
            "is_prompts": self.is_prompts,
            "use_cuda_kernel": self.use_cuda_kernel,
            "use_cuda_graph": self.use_cuda_graph,
            "kv_seq_len": self.kv_seq_len,
            "head_dim": self.head_dim,
            "high_precision": self.high_precision,
            "dtype": str(self.dtype).split(".")[-1],
            "use_spec_dec": self.use_spec_dec,
            "num_tokens_to_verify": self.num_tokens_to_verify,
            "batch_token_ids": self.batch_token_ids,
        }

    @staticmethod
    def from_rpc_param(rpc_dict: Dict[str, any]) -> "InputMetaData":
        """
        We intentionally don't use `dict.get` method to ensure we pass the right rpc param, or program will show error message