Functionals#

class tensorrt_llm.functional.AllReduceFusionOp( value, names=<not given>, *values, module=None, qualname=None, type=None, start=1, boundary=None, )[source]#

基类: IntEnum

LAST_PROCESS_FOR_UB = 2#

MOE_ALLREDUCE_RESIDUAL_RMS_NORM = 8#

NONE = 0#

RESIDUAL_RMS_NORM = 1#

RESIDUAL_RMS_NORM_OUT_QUANT_FP8 = 6#

RESIDUAL_RMS_NORM_OUT_QUANT_NVFP4 = 7#

RESIDUAL_RMS_NORM_QUANT_FP8 = 4#

RESIDUAL_RMS_NORM_QUANT_NVFP4 = 5#

RESIDUAL_RMS_PREPOST_NORM = 3#

class tensorrt_llm.functional.AllReduceParams( strategy: AllReduceStrategy = AllReduceStrategy.AUTO, fusion_op: AllReduceFusionOp = AllReduceFusionOp.NONE, bias: Tensor | None = None, residual: Tensor | None = None, norm_weight: Tensor | None = None, scale: Tensor | None = None, norm_pre_residual_weight: Tensor | None = None, eps: float = 1e-06, enable_allreduce: bool = True, )[source]#

基类: object

has_affine()[source]#

has_bias()[source]#

has_scale()[source]#

update_strategy()[source]#

class tensorrt_llm.functional.AllReduceStrategy( value, names=<not given>, *values, module=None, qualname=None, type=None, start=1, boundary=None, )[source]#

基类: IntEnum

AUTO = 3#

MIN_LATENCY = 1#

NCCL = 0#

ONESHOT = 4#

TWOSHOT = 5#

UB = 2#

class tensorrt_llm.functional.AttentionMaskType( value, names=<not given>, *values, module=None, qualname=None, type=None, start=1, boundary=None, )[源代码]#

基类: IntEnum

bidirectional = 3#

bidirectionalglm = 4#

blocksparse = 5#

causal = 1#

custom_mask = 6#

padding = 0#

sliding_window_causal = 2#

class tensorrt_llm.functional.Conditional(condition: Tensor)[源代码]#

基类: object

添加一个有条件地执行两个代码路径/子图的操作。

用法

conditional = Conditional(condition)
input_1_ = conditional.add_input(input_1) … input_n_ = conditional.add_input(input_n)
使用 input_1_, …, input_n_ 构建图以获得 true_output_value 和 false_output_value
output = conditional.add_output(true_output_value, false_output_value)

add_input( input: Tensor, ) → Tensor[源代码]#

add_output( true_value: Tensor, false_value: Tensor, ) → Tensor[源代码]#

class tensorrt_llm.functional.DimRange( shape: List[int | List[int] | Tuple[int, int, int]], names: List[str], )[源代码]#

基类: object

一个 DimRange 对象在一个优化配置中存储一个张量所有维度的范围。例如，张量有 2 个维度。然后数据成员为

self.min = [dim 0 min, dim 1 min] self.opt = [dim 0 opt, dim 1 opt] self.max = [dim 0 max, dim 1 max]

对于静态维度，它具有 min==opt==max，因此 ctor 中的 shape 参数可以是一个整数

class tensorrt_llm.functional.LayerNormPositionType( value, names=<not given>, *values, module=None, qualname=None, type=None, start=1, boundary=None, )[源代码]#

基类: IntEnum

post_layernorm = 1#

pre_layernorm = 0#

class tensorrt_llm.functional.LayerNormType( value, names=<not given>, *values, module=None, qualname=None, type=None, start=1, boundary=None, )[源代码]#

基类: IntEnum

GroupNorm = 2#

LayerNorm = 0#

RmsNorm = 1#

class tensorrt_llm.functional.MLPType( value, names=<not given>, *values, module=None, qualname=None, type=None, start=1, boundary=None, )[源代码]#

基类: IntEnum

FusedGatedMLP = 2#

GatedMLP = 1#

MLP = 0#

class tensorrt_llm.functional.PositionEmbeddingType( value, names=<not given>, *values, module=None, qualname=None, type=None, start=1, boundary=None, )[源代码]#

基类: IntEnum

alibi = 4#

alibi_with_scale = 5#

chatglm = 7#

static choices() → List[str][source]#

deferred = 10#

static from_string(s)[source]#

is_alibi() → bool[source]#

is_deferred() → bool[source]#

is_mrope() → bool[source]#

is_rope() → bool[source]#

learned_absolute = 0#

long_rope = 3#

mrope = 9#

relative = 6#

rope_gpt_neox = 2#

rope_gptj = 1#

yarn = 8#

class tensorrt_llm.functional.RopeEmbeddingUtils[source]#

基类: object

static apply_llama3_scaling( inv_freqs: ndarray, rope_scaling_config: dict, )[source]#

static apply_rotary_pos_emb( tensor: Tensor, position_embedding: List[Tensor] = None, pos_emb_type: PositionEmbeddingType = PositionEmbeddingType.rope_gptj, ) → Tensor[source]#

static apply_rotary_pos_emb_chatglm( qkv, position_embedding, num_attention_heads, attention_head_size, max_position_embeddings, rotary_embedding_scale, remove_input_padding, ) → Tensor[source]#

static apply_rotary_pos_emb_cogvlm( qkv, position_embedding, num_attention_heads, attention_head_size, max_position_embeddings, rotary_embedding_scale, remove_input_padding, ) → Tensor[source]#

static create_fake_weight( dim: int, dtype=<class 'numpy.float16'>, )[source]#

static create_sinusoidal_positions( num_pos: int, dim: int, theta: float = 10000.0, dtype=<class 'numpy.float32'>, )[source]#

static create_sinusoidal_positions_for_attention_plugin( num_pos: int, dim: int, theta: float = 10000.0, scale: float = 1.0, scale_type: ~tensorrt_llm.functional.RotaryScalingType = RotaryScalingType.none, rope_scaling_config: dict = None, dtype=<class 'numpy.float32'>, )[source]#

static create_sinusoidal_positions_for_cogvlm_attention_plugin( num_pos: int, dim: int, theta: float = 10000.0, scale: float = 1.0, scale_type: ~tensorrt_llm.functional.RotaryScalingType = RotaryScalingType.none, vision_start: int = 1, vision_length: int = 1225, dtype=<class 'numpy.float32'>, )[source]#

create_sinusoidal_positions_long_rope( num_orig_pos: int, dim: int, theta: float = 10000.0, scaling_short_factors: ~tensorrt_llm.functional.Tensor = 1.0, scaling_long_factors: ~tensorrt_llm.functional.Tensor = 1.0, short_mscale=None, long_mscale=None, dtype=<class 'numpy.float32'>, )[source]#

static create_sinusoidal_positions_yarn( num_pos: int, dim: int, base: int = 10000, scaling_factor: float = 1.0, original_max_position_embeddings: int = 4096, beta_fast: int = 32, beta_slow: int = 1, mscale: float = 1.0, mscale_all_dim: float = 1.0, dtype=<class 'numpy.float32'>, )[source]#

static rotate_every_two( tensor: Tensor, ) → Tensor[source]#

static rotate_half( tensor: Tensor, ) → Tensor[source]#

class tensorrt_llm.functional.RotaryScalingType( value, names=<not given>, *values, module=None, qualname=None, type=None, start=1, boundary=None, )[source]#

基类: IntEnum

dynamic = 2#

static from_string(s)[source]#

linear = 1#

llama3 = 4#

longrope = 3#

mrope = 6#

none = 0#

yarn = 5#

class tensorrt_llm.functional.SideStreamIDType( value, names=<not given>, *values, module=None, qualname=None, type=None, start=1, boundary=None, )[source]#

基类: IntEnum

disable = 0#

moe = 1#

class tensorrt_llm.functional.SliceInputType( value, names=<not given>, *values, module=None, qualname=None, type=None, start=1, boundary=None, )[source]#

基类: IntEnum

axes = 5#

data = 0#

fill_value = 4#

size = 2#

start = 1#

stride = 3#

class tensorrt_llm.functional.Tensor( name=None, dtype=None, shape=None, dim_range=None, is_network_input=True, location=<TensorLocation.DEVICE: 0>, network=None, trt_tensor=None, )[source]#

基类: object

表示稠密张量的类。

一个稠密张量是有名称的，具有形状，并且包含类型化的元素。张量的每个维度可以是静态的或动态的。静态维度在 TensorRT 引擎编译时是已知的。动态维度可以在运行时采用确定的值。张量可以位于主机 (CPU) 或设备 (GPU) 上。

abs()[source]#: 请参阅 functional.abs。

cast(dtype)[source]#: 请参阅 functional.cast。

property dtype#: 张量中元素的类型。

flatten(start_dim=0, end_dim=-1)[source]#: 请参阅 functional.flatten。

get_parent()[source]#: 获取生成此张量的层。

get_users()[source]#: 获取将此张量用作输入的层。

is_dynamic(dim=None)[source]#: 如果参数 ‘dim’ 为 None，则该函数返回一个布尔值，指示张量是否包含动态维度（True）或不包含（False）。在这种情况下，第一个维度被排除（因为它通常对应于批处理大小）。如果参数是一个整数，则该函数返回一个布尔值，指示维度 ‘dim’ 是否是动态的（True）或不是动态的（False）。

is_trt_wrapper()[source]#: 检查内部是否存在 trt.ITensor 成员，这是图重写器所必需的。为了区分用法，可能需要一个继承层次结构。

property location#: 张量的物理位置（在主机上或设备上）。

log()[source]#: 请参阅 functional.log。

mark_output( name: str | None = None, dtype: str | DataType | None = None, )[source]#

将张量标记为网络输出。

当一个张量被标记为输出时，它的内容可以在 TensorRT 引擎执行后获得。用户负责分配缓冲区来存储输出张量，以便准备执行 TensorRT 引擎。

max(dim, keepdim=False)[source]#: 请参阅 functional.max。

mean(dim, keepdim=False)[source]#: 请参阅 functional.mean。

property name#: 张量的名称。

ndim()[source]#: 返回张量的秩（即维度数）。

property network#

permute(dims)[source]#: 请参阅 functional.permute。

rank()[source]#: 返回张量的秩（即维度数）。

repeat(sizes)[source]#: 请参阅 functional.repeat

replace_all_uses_with(new_tensor)[source]#: 将此张量作为输入替换为消费者层的所有用途

select(dim, index)[source]#: 请参阅 functional.select。

property shape#: 张量的形状。

size(dim=None)[source]#: 如果 dim 参数为 None，则返回张量的形状。否则，返回 dim 指示的维度的大小。如果 dim 为负数或超出张量的秩，则该行为未定义。

split(split_size_or_sections, dim=0)[source]#: 请参阅 functional.split。

sqrt()[source]#: 请参阅 functional.sqrt。

squeeze(dim, zero_is_placeholder)[source]#: 请参阅 functional.squeeze。

transpose(dim0, dim1)[source]#: 请参阅 functional.transpose。

unbind(dim=0)[source]#: 请参阅 functional.unbind。

unsqueeze(dim)[source]#: 请参阅 functional.squeeze。

view(shape, zero_is_placeholder=True)[source]#: 请参阅 functional.view。

tensorrt_llm.functional.abs( input: ~tensorrt_llm.functional.Tensor, *, op: ~tensorrt.tensorrt.UnaryOperation = <UnaryOperation.ABS: 4>, ) → Tensor#

对单个输入添加逐元素运算。

以下闭包在 functional.* 中定义。

round for op=trt.UnaryOperation.ROUND sqrt for op=trt.UnaryOperation.SQRT exp for op=trt.UnaryOperation.EXP sin for op=trt.UnaryOperation.SIN cos for op=trt.UnaryOperation.COS abs for op=trt.UnaryOperation.ABS log for op=trt.UnaryOperation.LOG

它使用 TensorRT 中的 IUnaryLayer 实现。

参数:

input – Tensor 输入张量。
op – trt.UnaryOperation 要执行的一元运算。

返回值:

此逐元素运算生成的张量。

tensorrt_llm.functional.activation( input: Tensor, act_type: ActivationType, ) → Tensor[source]#

添加激活函数。

参数:

input – Tensor 应用激活函数的输入张量。
act_type – trt.ActivationType 激活的类型（RELU、TANH、SIGMOID、…）。

以下闭包在 functional.* 中定义。

relu for op=trt.ActivationType.RELU tanh for op=trt.ActivationType.TANH sigmoid for op=trt.ActivationType.SIGMOID

返回值:: 激活层生成的张量。

tensorrt_llm.functional.add( left: ~tensorrt_llm.functional.Tensor | int | float, right: ~tensorrt_llm.functional.Tensor | int | float, *, op: ~tensorrt.tensorrt.ElementWiseOperation = <ElementWiseOperation.SUM: 0>, ) → Tensor#

添加具有两个输入的逐元素运算。

对于每个输入，如果输入是整数或浮点数，该函数首先创建一个常量张量。然后，如果需要，它会扩展较小的张量以确保其秩与较大的张量相同。然后，它执行逐元素运算“op”。

以下闭包在 functional.* 中定义。

add for op=trt.ElementWiseOperation.SUM sub for op=trt.ElementWiseOperation.SUB mul for op=trt.ElementWiseOperation.PROD div for op=trt.ElementWiseOperation.DIV floordiv for op=trt.ElementWiseOperation.FLOOR_DIV gt for op=trt.ElementWiseOperation.GREATER lt for op=trt.ElementWiseOperation.LESS op_and for op=trt.ElementWiseOperation.AND op_or for op=trt.ElementWiseOperation.OR eq for op=trt.ElementWiseOperation.EQUAL minimum for op=trt.ElementWiseOperation.MIN maximum for op=trt.ElementWiseOperation.MAX pow for op=trt.ElementWiseOperation.POW

它使用 TensorRT 中的 IElementWiseLayer 实现。

参数:

left – Union[Tensor, int, float] 第一个输入。如果该输入是整数或浮点数，则该函数会创建一个常量张量。
right – Union[Tensor, int, float] 第二个输入。如果该输入是整数或浮点数，则该函数会创建一个常量张量。
op – trt.ElementWiseOperation 要执行的二元运算。

返回值:

此逐元素运算生成的张量。

tensorrt_llm.functional.allgather( tensor: Tensor, group: List[int], gather_dim: int = 0, ) → Tensor[source]#

添加执行集体 all-gather 的运算。

让我们将“group_size”定义为“group”列表的长度。该函数创建一个层来收集分布在“group_size”个参与等级（每个等级一个 GPU）中的“group_size”个张量。

列表“group”包含参与集体操作的等级的标识符。

请注意，这里的“group”可以是 TP 组或 PP 组，因为 allgather 通信不限于特定的拆分模式。因此，“group_size”不需要等于 MPI “world_size”。

不同等级中的张量必须是 1D 张量（或视图），并且输出张量将具有相同的形状。

给定“section_size = input.shape[0] / group_size”，每个等级都会贡献其输入张量的一部分，该部分对应于“rank*section_size:(rank+1)*section_size”。

该操作使用一个插件来实现，该插件包装了 NCCL all-gather 集体操作。有关详细信息，请参见 https://docs.nvda.net.cn/deeplearning/nccl/user-guide/docs/usage/collectives.html#allgather。

参数:

tensor – Tensor 输入张量。
group – List[int] 参与 all-gather 操作的等级。
gather_dim – int = 0 沿给定维度收集。默认情况下为 0，即视为 1D 张量。

返回值:

该层生成的张量。

tensorrt_llm.functional.allreduce(tensor: ~tensorrt_llm.functional.Tensor, group: ~typing.List[int], all_reduce_params: ~tensorrt_llm.functional.AllReduceParams | None = <tensorrt_llm.functional.AllReduceParams object>) → Tensor[source]#

添加执行集体 all-reduce 的操作。

让我们将“world_size”定义为“group”列表的长度。该函数创建一个层来计算分布在“world_size”个参与等级（每个等级一个 GPU）中的“world_size”个张量的总和。

列表“group”包含参与集体操作的等级的标识符。

不同等级中的张量必须是 1D 张量（或视图），并且输出张量将具有相同的形状。输出张量将在“world_size”个等级上复制。

该操作使用一个插件来实现，该插件包装了 NCCL all-reduce 集体操作。有关详细信息，请参见 https://docs.nvda.net.cn/deeplearning/nccl/user-guide/docs/usage/collectives.html#allreduce。

参数:

tensor – Tensor 输入张量。
group – List[int] 参与 all-reduce 操作的等级。
strategy – AllReduceStrategy NCCL 将 all-reduce 委托给 NCCL，而 ONESHOT 和 TWOSHOT 是自定义的延迟优化算法。 AUTO 根据消息大小启发式算法在三者之间进行选择。

返回值:

该层生成的张量。

tensorrt_llm.functional.arange( start: Tensor | int, end: Tensor | int, dtype: str, ) → Tensor[source]#

添加用于填充 1D 张量的操作。

张量填充的值介于开始和结束之间，不同元素之间的步长为 1。在伪代码中，它对应于使用以下值填充的张量

output = Tensor([dtype(ii) for ii in range(start, end, 1)])

例如，调用 arange(3, 6, ‘int32’) 将向 TensorRT 图添加一个操作，该操作在执行时将生成 [3, 4, 5]。调用 arange(2, 5, ‘float32’) 将添加一个层以生成 [2.0, 3.0, 4.0]。

此操作使用 trt.FillOperation.LINSPACE 模式下的 tensorrt.IFillLayer 实现。

参数:

start – Union[Tensor, int] 范围的起始点。
end – Union[Tensor, int] 范围的终点。
dtype – str 元素的类型。有关支持的类型和类型名称的列表，请参见 _utils.py 中的 _str_to_trt_dtype_dict。

返回值:

填充层生成的张量。它是包含 end-start 个类型为 dtype 的元素的 1D 张量。

tensorrt_llm.functional.argmax( input: Tensor, dim: int, keepdim: bool = False, ) → Tensor[source]#

添加一个 argmax 运算。

正如 ONNX 文档中所解释的，

onnx/onnx

该函数创建一个层，用于计算输入张量的元素沿提供的 dim 的最大元素的索引。如果 keepdims 为 True，则结果张量与输入张量具有相同的秩。如果 keepdims 为 False，则结果张量会修剪缩减的维度。

参数:

input – Tensor 输入张量。
dim – int 计算 argmax 索引的维度。
keepdim – bool 我们是否保留执行缩减的维度？如果设置为 True，则为是，否则为否。

返回值:

此 argmax 运算生成的张量。

tensorrt_llm.functional.assertion( condition: Tensor, message: str = '', ) → None[源代码]#

tensorrt_llm.functional.avg_pool2d( input: Tensor, kernel_size: Tuple[int], stride: Tuple[int] | None = None, padding: Tuple[int] | None = (0, 0), ceil_mode: bool = False, count_include_pad: bool = True, ) → Tensor[源代码]#

tensorrt_llm.functional.bert_attention( tensor: Tensor, input_lengths: Tensor, num_heads: int, head_size: int, q_scaling: float, relative_attention: bool = False, relative_attention_bias: Tensor = None, max_distance: int = 0, max_input_length: Tensor = None, sage_attn: bool = False, sage_attn_q_block_size: int = 0, sage_attn_k_block_size: int = 0, sage_attn_v_block_size: int = 0, ) → Tuple[Tensor][源代码]#

添加一个在 BERT 中执行多头注意力机制的操作。

多头注意力 (MHA) 是一个批量矩阵乘法、一个 softmax 和一个批量矩阵乘法的序列，如 https://arxiv.org/abs/1706.03762 中所述。该函数添加一个操作，使用单个 GPU 内核执行这些计算。

输入张量包含 Q、K 和 V 元素。它是一个 2D 张量，其形状为“[sum_of_tokens, 3*hidden_dim]”，其中“sum_of_tokens”是批次中序列长度的总和。

在 MHA 中，Q*K^T 乘积的输出按一个常数值缩放，该常数值计算为

1.f / (q_scaling * sqrt(head_size)).

该“q_scaling”常数是该函数的最后一个参数。

该层是使用插件实现的（参见 bertAttentionPlugin）。

参数:

tensor – Tensor QKV 输入张量。
input_lengths – Tensor 每个序列的长度。它是一个大小为“batch_size”的 1D 张量。
num_heads – int 头数。
head_size – int 每个头的大小。
q_scaling – float 用于计算缩放因子的因子，以缩放“Q*K^T”乘积的输出。
relative_attention – bool = False 如果启用相对注意力。
relative_attention_bias – Tensor = None 相对注意力偏差 [num_heads, max_seq_len, max_seq_len]，或隐式模式的相对注意力嵌入表 [num_heads, num_buckets]。
max_distance – int = 0 注意力中相对位置的最大距离，用于隐式模式。默认值为 0，表示使用相对注意力偏差的常规模式。仅当传入非零的正 max_distance 值时，才会启用隐式模式。请参阅 docs/source/advanced/gpt-attention.md 中的相对注意力偏差
max_input_length – Tensor = None 张量形状表示的最大输入序列长度。需要 remove_input_padding 预定义插件工作空间大小。
sage_attn – bool = False SageAttention 是 attention 内核的 8 位实现。它的输入 q、k、v 和输出数据类型为 16 位。它每次在 attention 之前对 q、k、v 张量执行动态量化。 thu-ml/SageAttention
sage_attn_q_quant_size – int = 0 沿 q 张量序列维度的动态量化块大小。每个量化块将共享一个比例。
sage_attn_k_quant_size – int = 0 沿 k 张量序列维度的动态量化块大小。每个量化块将共享一个比例。
sage_attn_v_quant_size – int = 0 沿 v 张量序列维度的动态量化块大小。每个量化块将共享一个比例。

返回值:

该层生成的张量。

tensorrt_llm.functional.broadcast_helper( left: Tensor | int | float, right: Tensor | int | float, ) → Tuple[Tensor, Tensor][源代码]#

用于执行广播的辅助函数。

对于每个输入，如果输入是整数或浮点数，则该函数首先创建一个常量张量。然后，如果需要，它会扩展较小的张量，以确保其秩与较大的张量相同。

参数:

left – Union[Tensor, int, float] 第一个输入。如果该输入是整数或浮点数，则该函数会创建一个常量张量。
right – Union[Tensor, int, float] 第二个输入。如果该输入是整数或浮点数，则该函数会创建一个常量张量。

返回值:

一对相同秩的张量。

tensorrt_llm.functional.cast( input: Tensor, dtype: str | DataType, )[源代码]#

添加类型转换操作。

对于 INT8 类型的输入张量，此函数将输入的动态范围设置为 [-127, 127]，以便自动反量化。对于强制转换为 INT8，该函数将输出的动态范围设置为 [-127, 127]，以便自动量化。

参数:

input – Tensor 在其上应用强制转换的输入张量。
dtype – str 或 trt.DataType 强制转换后输出张量的数据类型。当“dtype”作为字符串提供时，它必须是有效名称中的名称。有关受支持类型和类型名称的列表，请参见 _utils.py 中的 _str_to_trt_dtype_dict。

返回值:

插入层生成的张量。

tensorrt_llm.functional.categorical_sample( probs: Tensor, rand_data: Tensor = None, ) → Tensor[源代码]#

这是一个采样操作，等效于 torch.distributions.Categorical.sample()，即给定一个概率分布张量，它对该张量的索引进行采样。参见： https://pytorch.ac.cn/docs/stable/distributions.html#torch.distributions.categorical.Categorical.sample 注意：这假设给定的概率**未**归一化。

参数:

probs – Tensor 一个 1-D 浮点张量，表示概率分布。
rand_data – Tensor (可选) 形状与 probs 张量相同的随机张量。如果未提供，此函数将添加一个 rand() 运算来生成它并用于采样。

返回值:

一个包含 probs 张量的单个索引的张量，表示样本。

tensorrt_llm.functional.chunk( tensor: Tensor, chunks: int, dim: int = 0, ) → Tensor[源代码]#

添加将张量分割成子张量的操作。

此操作创建一个张量列表，这些张量通过沿维度 ‘dim’ 对输入张量进行分块而获得。它产生 ‘chunks’ 个子张量。

该操作仅针对静态张量（无动态维度）定义，并且张量在维度 ‘dim’ 中的大小必须是 ‘chunks’ 的倍数：‘input.shape[dim] % chunks == 0’。

它映射到 ‘split’，其中 ‘split_size = input.shape[dim] / chunks’。

参数:

tensor – Tensor 要切片的输入张量。
chunks – int 要将输入张量拆分成的切片数。
dim – int 要切片的张量的维度。

返回值:

由不同操作产生的张量列表。

tensorrt_llm.functional.clip( input: Tensor, alpha: float, beta: float, ) → Tensor[源代码]#

添加一个 CLIP 操作，将范围设置为 [alpha, beta]。

参数:

input – Tensor 应用激活函数的输入张量。
alpha – float CLIP 函数的下限。
beta – float CLIP 函数的上限。

返回值:

激活层生成的张量。

tensorrt_llm.functional.concat( inputs: Sequence[Tensor | int], dim: int = 0, ) → Tensor[源代码]#

添加一个操作来连接张量。

该函数创建一个操作，该操作连接来自序列 ‘inputs’ 的张量。连接沿维度 ‘dim’ 完成。

‘inputs’ 中的所有张量必须具有相同的形状，除了维度 ‘dim’。

for ii in range(inputs[0].rank())
assert (ii == dim) or all(inp.shape[ii] == inputs[0].shape[ii] for inp in inputs)

输出张量的形状定义为

for ii in range(inputs[0].rank())
# 与维度 ii != dim 中的所有输入大小相同。 output.shape[ii] = inputs[0].shape[ii]

# 不同输入在维度 ‘dim’ 中的大小之和。 if ii == dim

for jj in range(1, len(inputs))
output.shape[ii] += inputs[jj].shape[ii]

例如，给定一个由两个 2D 张量 [[0, 1], [2, 3]] 和 [[4, 5], [6, 7]] 组成的序列，它们的形状均为 [2, 2]，

concat(inputs, 0)

将产生 [[0, 1], [2, 3], [4, 5], [6, 7]]，形状为 [4, 2]，并且

concat(inputs, 1)

将产生 [[0, 1, 4, 5], [2, 3, 6, 7]]，形状为 [2, 4]。

参数:

inputs – Sequence[Union[Tensor, int]] 要连接的张量序列。对于整数，该函数创建常量张量。
dim – int 执行连接的维度。

返回值:

包含张量连接的张量。

tensorrt_llm.functional.constant( ndarray: ndarray, as_dtype: DataType | None = None, as_shape=None, ) → Tensor[源代码]#

添加一个常量层。

TensorRT 图以常量层（tensorrt.IConstantLayer）的形式封装常量值。该函数从一个 Numpy 值数组创建一个这样的层。在 TensorRT 编译网络后，这些权重存储在序列化的 TensorRT 引擎中。

参数:: ndarray – numpy.ndarray 由该常量层封装的值（权重）数组。
返回值:: 插入层生成的张量。

tensorrt_llm.functional.constant_to_tensor_( input: Tensor | int | float | bool, dtype: DataType | str = None, to_array=True, ) → Tensor[源代码]#

tensorrt_llm.functional.constants_to_tensors_( *inputs: Tensor | int | float, ) → Tuple[Tensor, ...][源代码]#

用于从多个输入创建张量的辅助函数。

对于每个输入，如果输入是整数或浮点数，该函数首先创建一个常量张量。然后，如果任何输入是 int64，它会将其他整数输入向上转换为 int64。

参数:: inputs – Tuple[Union[Tensor, int, float], …] 用于创建张量的输入。
返回值:: 张量的元组。

tensorrt_llm.functional.conv1d( input: Tensor, weight: Tensor, bias: Tensor | None = None, stride: int = 1, padding: int = 0, dilation: int = 1, groups: int = 1, ) → Tensor[源代码]#

tensorrt_llm.functional.conv2d( input: Tensor, weight: Tensor, bias: Tensor | None = None, stride: Tuple[int, int] = (1, 1), padding: Tuple[int, int] = (0, 0), dilation: Tuple[int, int] = (1, 1), groups: int = 1, pre_padding: Tuple[int, int] | None = None, post_padding: Tuple[int, int] | None = None, ) → Tensor[源代码]#

tensorrt_llm.functional.conv3d( input: Tensor, weight: Tensor, bias: Tensor | None = None, stride: int | Tuple[int, int] = (1, 1, 1), padding: int | Tuple[int, int] = (0, 0, 0), dilation: int | Tuple[int, int] = (1, 1, 1), groups: int = 1, ) → Tensor[源代码]#

tensorrt_llm.functional.conv_transpose2d( input: Tensor, weight: Tensor, bias: Tensor | None = None, stride: Tuple[int, int] = (1, 1), padding: Tuple[int, int] = (0, 0), output_padding: Tuple[int, int] = (0, 0), dilation: Tuple[int, int] = (1, 1), groups: int = 1, ) → Tensor[源代码]#

tensorrt_llm.functional.cos( input: ~tensorrt_llm.functional.Tensor, *, op: ~tensorrt.tensorrt.UnaryOperation = <UnaryOperation.COS: 7>, ) → Tensor#

对单个输入添加逐元素运算。

以下闭包在 functional.* 中定义。

round for op=trt.UnaryOperation.ROUND sqrt for op=trt.UnaryOperation.SQRT exp for op=trt.UnaryOperation.EXP sin for op=trt.UnaryOperation.SIN cos for op=trt.UnaryOperation.COS abs for op=trt.UnaryOperation.ABS log for op=trt.UnaryOperation.LOG

它使用 TensorRT 中的 IUnaryLayer 实现。

参数:

input – Tensor 输入张量。
op – trt.UnaryOperation 要执行的一元运算。

返回值:

此逐元素运算生成的张量。

tensorrt_llm.functional.cp_split_plugin( input_ids: Tensor, host_request_types: Tensor, host_context_lengths: Tensor, cp_size: int = 1, cp_rank: int = 0, ) → Tensor[源代码]#

添加一个操作来执行上下文并行拆分。

此操作将 input_ids 拆分为 cp_size 个块，并返回第 cp_rank 个块。当 seqlen % cp_size != 0 时，每个 rank 的块大小将为 [seqlen // cp_size, seqlen // cp_size, …, seqlen - (seqlen // cp_size) * cp_size]

它插入一个 IPluginV3Layer。

参数:

input – Tensor 输入张量包含要拆分的索引。
host_request_types – Tensor = None (在 CPU 上) 主机上的张量，指示请求是处于上下文阶段还是生成阶段。它的形状是 [batch_size]。请参阅 docs/gpt_attention.md 中的 Inflight Batching，
host_context_lengths – Tensor = None (在 CPU 上) 一个主机张量，包含不同输入的长度

返回值:

输出的拆分张量。输出的拆分张量的长度。重建序列的索引

tensorrt_llm.functional.create_allreduce_plugin( network: INetworkDefinition, tensor: ITensor, workspace: ITensor | None, group: array, dtype: DataType, all_reduce_params: AllReduceParams, )[源代码]#

tensorrt_llm.functional.cuda_stream_sync( input_list: List[Tensor], side_stream_id: SideStreamIDType, ) → Tensor[源代码]#

在主流上等待侧流。 output = input_list[0]

参数:

input_list – List[Tensor] (在 GPU 上) 输入张量列表。
side_stream_id – int (在 CPU 上) 侧流 ID。

tensorrt_llm.functional.cumsum( input: Tensor, dim: int, prefer_plugin: bool = True, ) → Tensor[源代码]#

添加一个操作，用于计算给定维度中张量元素的包含累积和。

给定一个输入张量，该函数创建一个操作，计算维度 ‘dim’ 中元素的包含累积和，以创建一个新的张量。输出张量与输入张量具有相同的形状。

输入张量的秩必须 >= 1。“dim”必须有效，并且支持负值。

例如，对于输入=[[4, 2, 5], [2, 1, 2], [4, 7, 1]]，其形状为 [3, 3]，

cumsum(input, 0)

将产生 [[4, 2, 5], [6, 3, 7], [10, 10, 8]]。

cumsum(input, 1)

将产生 [[4, 6, 11], [2, 3, 5], [4, 11, 12]]。

该操作由 TensorRT ILoopLayer 实现。

参数:

input – Tensor 用于计算包含累积和的输入张量。
dim – int 用于计算包含累积和的维度。支持负值。
prefer_plugin – bool 如果 dim 是最后一个维度，是否使用 cumsumLastDim 插件。

返回值:

包含输入包含累积和的张量。

tensorrt_llm.functional.div( left: ~tensorrt_llm.functional.Tensor | int | float, right: ~tensorrt_llm.functional.Tensor | int | float, *, op: ~tensorrt.tensorrt.ElementWiseOperation = <ElementWiseOperation.DIV: 5>, ) → Tensor#

添加具有两个输入的逐元素运算。

以下闭包在 functional.* 中定义。

add for op=trt.ElementWiseOperation.SUM sub for op=trt.ElementWiseOperation.SUB mul for op=trt.ElementWiseOperation.PROD div for op=trt.ElementWiseOperation.DIV floordiv for op=trt.ElementWiseOperation.FLOOR_DIV gt for op=trt.ElementWiseOperation.GREATER lt for op=trt.ElementWiseOperation.LESS op_and for op=trt.ElementWiseOperation.AND op_or for op=trt.ElementWiseOperation.OR eq for op=trt.ElementWiseOperation.EQUAL minimum for op=trt.ElementWiseOperation.MIN maximum for op=trt.ElementWiseOperation.MAX pow for op=trt.ElementWiseOperation.POW

它使用 TensorRT 中的 IElementWiseLayer 实现。

参数:

left – Union[Tensor, int, float] 第一个输入。如果该输入是整数或浮点数，则该函数会创建一个常量张量。
right – Union[Tensor, int, float] 第二个输入。如果该输入是整数或浮点数，则该函数会创建一个常量张量。
op – trt.ElementWiseOperation 要执行的二元运算。

返回值:

此逐元素运算生成的张量。

tensorrt_llm.functional.dora_plugin( activations: Tensor, out_hidden_sizes: list[int], lora_weights_pointers: list[Tensor], host_request_types: Tensor, host_context_lengths: Tensor | None = None, ) → Tensor[source]#

DoRA 插件将列向缩放应用于 LoRA 层的输出。

参数:

input – Tensor (在 GPU 上) 输入张量。其形状为 [batch_size, seq_len, dim] 或 [num_tokens, dim]（用于 remove_input_padding）
out_hidden_sizes – list[int] 相关 LoRA 模块中每个适配器的输出隐藏大小。例如，对于 qkv 投影，out_hidden_sizes 应为 [q_dim, k_dim, v_dim]。
host_request_types – Tensor = None 主机上的张量，指示请求处于上下文阶段还是生成阶段。其形状为 [batch_size]。请参阅 docs/source/advanced/gpt-attention.md 中的 Inflight Batching。
host_context_lengths – cpu Tensor = None 包含不同输入长度的主机张量。

返回值:

该层生成的张量。

tensorrt_llm.functional.einsum( einsum_eq: str, inputs: Sequence[Tensor], ) → Tensor[source]#

添加 Einsum 操作。

该操作映射到 tensorrt.IEinsumLayer。正如 TensorRT 文档中所述，该层根据爱因斯坦求和约定，沿着由 equation 参数指定的维度对输入的元素进行求和。该层可以有一个或多个秩 >= 0 的输入。所有输入必须具有相同的数据类型。该层支持除 bool 之外的所有 TensorRT 数据类型。有一个输出张量，其类型与输入张量相同。输出张量的形状由 equation 确定。

equation 为每个输入中的每个维度指定 ASCII 小写字母，维度顺序与输入顺序相同，并用逗号分隔每个输入。用相同下标标记的维度必须匹配或能够广播。一个输入中重复的下标标签取对角线。跨多个输入重复一个标签表示这些轴将被相乘。从输出中省略标签表示将对沿着这些轴的值进行求和。在隐式模式下，表达式中出现一次的索引将按字母顺序递增成为输出的一部分。在显式模式下，可以通过添加箭头 (‘->’) 后跟输出的下标来控制输出。例如，“ij,jk->ik” 等同于 “ij,jk”。省略号 (‘…’) 可以代替下标来广播维度。有关 equation 语法的更多详细信息，请参阅 TensorRT 开发者指南。

许多常见操作可以使用 Einsum equation 来表达。对于 .. rubric:: 示例

矩阵转置：ij->ji 求和：ij-> 矩阵-矩阵乘法：ik,kj->ij 点积：i,i-> 矩阵-向量乘法：ik,k->i 批量矩阵乘法：ijk,ikl->ijl 批量对角线：…ii->…i

请注意，TensorRT 不支持省略号或对角线操作，因此 TensorRT-LLM 也不支持。

参数:

einsum_eq – str Einsum equation。
inputs – Sequence[Tensor] Einsum 操作使用的输入序列。

返回值:

Einsum 操作生成的张量。

tensorrt_llm.functional.elementwise_binary( left: Tensor | int | float, right: Tensor | int | float, op: ElementWiseOperation, ) → Tensor[source]#

添加具有两个输入的逐元素运算。

以下闭包在 functional.* 中定义。

add for op=trt.ElementWiseOperation.SUM sub for op=trt.ElementWiseOperation.SUB mul for op=trt.ElementWiseOperation.PROD div for op=trt.ElementWiseOperation.DIV floordiv for op=trt.ElementWiseOperation.FLOOR_DIV gt for op=trt.ElementWiseOperation.GREATER lt for op=trt.ElementWiseOperation.LESS op_and for op=trt.ElementWiseOperation.AND op_or for op=trt.ElementWiseOperation.OR eq for op=trt.ElementWiseOperation.EQUAL minimum for op=trt.ElementWiseOperation.MIN maximum for op=trt.ElementWiseOperation.MAX pow for op=trt.ElementWiseOperation.POW

它使用 TensorRT 中的 IElementWiseLayer 实现。

参数:

left – Union[Tensor, int, float] 第一个输入。如果该输入是整数或浮点数，则该函数会创建一个常量张量。
right – Union[Tensor, int, float] 第二个输入。如果该输入是整数或浮点数，则该函数会创建一个常量张量。
op – trt.ElementWiseOperation 要执行的二元运算。

返回值:

此逐元素运算生成的张量。

tensorrt_llm.functional.embedding( input: Tensor, weight: Tensor, tp_size=1, tp_group=None, sharding_dim=0, tp_rank=None, per_token_scale=None, padding=None, ) → Tensor[source]#

添加执行嵌入查找的操作。

该操作执行嵌入查找。 ‘input’ 张量包含要收集的 ‘weight’ 行的标识符。

1. 在多个 GPU 上分配嵌入查找表当 ‘tp_size’ 大于 1 且定义了 ‘tp_group’ 时，此嵌入查找将在多个 GPU 之间分配。

当 ‘sharding_dim==0’ 时，每个 GPU 存储嵌入表行的子集（每个 GPU 的行数由 weights.shape[0] 给出，并且 GPU 上存储的第一行的偏移量由 rank * weights.shape[0] 给出）。每个并行 rank 将查询所有索引，并为未存储在关联 GPU 上的权重设置 0。为了计算最终结果，将一个并行的 all-reduce 操作添加到 TensorRT 图中。可以使用插件或 TensorRT 支持的运算符来执行该查找。

当 ‘sharding_dim==1’ 时，每个 GPU 存储嵌入表列的子集。每个 rank 都可以获得一部分嵌入结果。然后，使用 all-gather 操作来收集嵌入。相关的转置操作也用于获得最终结果。

2. 将嵌入查找表作为一个整体存储当 ‘tp_size’ 不大于 1 时，将不会分割嵌入查找表。在这种情况下，当设置了 default_net().plugin_config.lookup_plugin 时，使用插件来实现该操作（没有 all-reduce 操作）。否则，将使用 TensorRT 中的标准 IGatherLayer 来实现此操作。

参数:

input – Tensor 包含执行查找的索引的输入张量。
weight – Tensor 要从中收集的表。
tp_size – int 协作执行该嵌入的 GPU 数量。
tg_group – Optional[List[int]] 当 tp_size > 1 时，参与 all-reduce 的 world rank 组。
sharding_dim – int sharding_dim = 0 表示我们在词汇维度中对嵌入表进行分片； sharding_dim = 1 表示我们在嵌入维度中对嵌入表进行分片。
tp_rank – int 张量并行 rank。用于计算词汇维度上的 TP 中的偏移量。
padding – Tensor 添加到嵌入表末尾的其他填充，然后再馈送到 gather 操作中。

返回值:

嵌入查找层生成的张量。

tensorrt_llm.functional.eq( left: ~tensorrt_llm.functional.Tensor | int | float, right: ~tensorrt_llm.functional.Tensor | int | float, *, op: ~tensorrt.tensorrt.ElementWiseOperation = <ElementWiseOperation.EQUAL: 11>, ) → Tensor#

添加具有两个输入的逐元素运算。

以下闭包在 functional.* 中定义。

add for op=trt.ElementWiseOperation.SUM sub for op=trt.ElementWiseOperation.SUB mul for op=trt.ElementWiseOperation.PROD div for op=trt.ElementWiseOperation.DIV floordiv for op=trt.ElementWiseOperation.FLOOR_DIV gt for op=trt.ElementWiseOperation.GREATER lt for op=trt.ElementWiseOperation.LESS op_and for op=trt.ElementWiseOperation.AND op_or for op=trt.ElementWiseOperation.OR eq for op=trt.ElementWiseOperation.EQUAL minimum for op=trt.ElementWiseOperation.MIN maximum for op=trt.ElementWiseOperation.MAX pow for op=trt.ElementWiseOperation.POW

它使用 TensorRT 中的 IElementWiseLayer 实现。

参数:

left – Union[Tensor, int, float] 第一个输入。如果该输入是整数或浮点数，则该函数会创建一个常量张量。
right – Union[Tensor, int, float] 第二个输入。如果该输入是整数或浮点数，则该函数会创建一个常量张量。
op – trt.ElementWiseOperation 要执行的二元运算。

返回值:

此逐元素运算生成的张量。

tensorrt_llm.functional.exp( input: ~tensorrt_llm.functional.Tensor, *, op: ~tensorrt.tensorrt.UnaryOperation = <UnaryOperation.EXP: 0>, ) → Tensor#

对单个输入添加逐元素运算。

以下闭包在 functional.* 中定义。

round for op=trt.UnaryOperation.ROUND sqrt for op=trt.UnaryOperation.SQRT exp for op=trt.UnaryOperation.EXP sin for op=trt.UnaryOperation.SIN cos for op=trt.UnaryOperation.COS abs for op=trt.UnaryOperation.ABS log for op=trt.UnaryOperation.LOG

它使用 TensorRT 中的 IUnaryLayer 实现。

参数:

input – Tensor 输入张量。
op – trt.UnaryOperation 要执行的一元运算。

返回值:

此逐元素运算生成的张量。

tensorrt_llm.functional.expand( input: Tensor, expand_shape: Tensor, ) → Tensor[source]#

添加一个操作来扩展张量。

该操作将输入张量在单例维度中扩展到 expand_shape 张量中相应维度所指示的大小。换句话说，给定一个维度大小为 1 的输入张量，这些维度将被扩展到 expand_shape 中的大小。

例如，形状为 [4, 3, 1, 3] 的张量将通过使用 expand(input, [4, 3, 2, 3]) 创建的层扩展为形状为 [4, 3, 2, 3] 的张量。

扩展可以复制值，也可以映射到扩展维度中步长为 0 的视图。例如，对于形状为 [1, 2] 的张量 [[3, 2]]，

expand([[3, 2]], [2, 2])

可用于将输入扩展到 [[3, 2], [3, 2]]。

此操作是使用 tensorrt.ISliceLayer 实现的。当前的实现不验证非单例维度是否缩小。换句话说，对于形状为 [4, 1, 2] 的输入，

expand(input, [3, 2, 2])

将生成形状为 [3, 2, 2] 的张量。该行为将来可能会更改。

参数:

input – Tensor 输入张量。
expand_shape – 张量扩展张量的新形状。

返回值:

expand 层生成的张量。

tensorrt_llm.functional.expand_dims( input: Tensor, dim: int | Sequence[int], shape_cast_dtype=None, ) → Tensor[source]#

添加一个操作以使用单例维度扩展张量形状。

该函数向网络添加一个 tensorrt.IShuffleLayer。给定一个秩为 N 的 ‘input’ 和 M 个维度的序列，此操作生成的输出张量（由 TensorRT 执行时）将具有 N+M 的秩。单例维度将被插入到 ‘dim’ 中的不同位置。

该操作的伪代码是

new_shape, ii = [], 0 for jj in range(input.rank() + len(dim))

new_shape.append(1 if jj in dims else input.shape[ii++])

例如，对于形状为 [3, 4, 1, 5] 的张量

expand_dims(input, [0, 2])

将生成形状为 [1, 3, 1, 4, 1, 5] 的张量。

参数:

input – 张量要扩展的输入张量。
dim – Union[int, Sequence[int]] 输出张量中插入单例维度的位置。

返回值:

shuffle 层生成的张量。

tensorrt_llm.functional.expand_dims_like( left: Tensor | int | float, right: Tensor, ) → Tensor[source]#

添加一个操作以将第一个张量扩展到与第二个张量相同的秩。

该函数接受第一个张量。它还接受整数或浮点数，在这种情况下，它会从中创建一个常量张量。在这两种情况下，都将第一个张量的秩与第二个张量的秩进行比较。如果它们的秩相同，则返回第一个张量。否则，第一个张量将在左侧扩展以匹配第二个张量的秩。

请注意，形状不必匹配，在该函数中仅考虑秩。

例如，对于形状为 [3, 4] 和 [4, 3, 2] 的一对张量，第一个张量将被扩展为秩为 3 且形状为 [1, 3, 4] 的张量。

参数:

left – Union[Tensor, int, float] 要扩展的第一个张量。当提供标量值作为参数时，该函数首先创建一个张量，然后再扩展它（如果需要）。
right – 张量要匹配的参考张量。

返回值:

shuffle 层生成的张量。

tensorrt_llm.functional.expand_mask( mask: Tensor, tgt_len: Tensor | None = None, ) → Tensor[source]#

扩展 attention mask。

该函数添加一系列操作，以从形状为“[batch_size, src_seq_len]”的张量扩展到形状为“[batch_size, 1, tgt_seq_len, src_seq_len]”的张量。它可用于创建应用于多头 attention 块中 softmax 运算之前的 Q*K^T 乘积的掩码。

参数:

mask – 张量输入掩码
tgt_len – Optional[Tensor] 输出张量中第三个维度的尺寸。如果为 None，则使用输入的第二个维度。

返回值:

由该系列操作创建的张量。

tensorrt_llm.functional.flatten( input: Tensor, start_dim: int = 0, end_dim: int = -1, )[source]#

通过将输入重塑为一维张量来展平输入。

如果传递了 start_dim 或 end_dim，则仅展平以 start_dim 开头并以 end_dim 结尾的维度。输入中元素的顺序不变。

参数:

input – 张量要展平的输入张量。
start_dim – int 要展平的第一个维度。
end_dim – int 要展平的最后一个维度。

返回值:

flatten 层生成的张量。

tensorrt_llm.functional.flip( input: Tensor, dims: Sequence[int], ) → Tensor[source]#

沿 dims 中给定的轴反转 n 维张量的顺序。

该 flip 操作映射到 TensorRT ISliceLayer。对于 dims 中列出的维度，它将元素从最后一个复制到第一个（从 (N-1) 向下到 0，步长为 -1）。对于不在 ‘dims’ 中的维度，它将元素从第一个复制到最后一个（从 0 到 N-1，步长为 1）。

参数:

input – Tensor 在其上应用强制转换的输入张量。
dims – list 或 tuple 要翻转的轴。支持负索引。

返回值:

插入层生成的张量。

tensorrt_llm.functional.floordiv( left: ~tensorrt_llm.functional.Tensor | int | float, right: ~tensorrt_llm.functional.Tensor | int | float, *, op: ~tensorrt.tensorrt.ElementWiseOperation = <ElementWiseOperation.FLOOR_DIV: 7>, ) → Tensor#

添加具有两个输入的逐元素运算。

以下闭包在 functional.* 中定义。

add for op=trt.ElementWiseOperation.SUM sub for op=trt.ElementWiseOperation.SUB mul for op=trt.ElementWiseOperation.PROD div for op=trt.ElementWiseOperation.DIV floordiv for op=trt.ElementWiseOperation.FLOOR_DIV gt for op=trt.ElementWiseOperation.GREATER lt for op=trt.ElementWiseOperation.LESS op_and for op=trt.ElementWiseOperation.AND op_or for op=trt.ElementWiseOperation.OR eq for op=trt.ElementWiseOperation.EQUAL minimum for op=trt.ElementWiseOperation.MIN maximum for op=trt.ElementWiseOperation.MAX pow for op=trt.ElementWiseOperation.POW

它使用 TensorRT 中的 IElementWiseLayer 实现。

参数:

left – Union[Tensor, int, float] 第一个输入。如果该输入是整数或浮点数，则该函数会创建一个常量张量。
right – Union[Tensor, int, float] 第二个输入。如果该输入是整数或浮点数，则该函数会创建一个常量张量。
op – trt.ElementWiseOperation 要执行的二元运算。

返回值:

此逐元素运算生成的张量。

tensorrt_llm.functional.gather( input: Tensor, dim: int, indices: Tensor | int, ) → Tensor[source]#

添加一个操作，从张量中收集元素。

该函数实现了 ONNX 规范中的 GatherElements 操作符，如下所述：

onnx/onnx

input 和 indices 参数必须具有相同的 rank >= 1。该操作将生成一个与 indices 张量具有相同形状的张量。axis 是要收集的维度。

如 ONNX 描述中所示，对于 3D 张量，输出为：

out[i][j][k] = input[indices[i][j][k]][j][k] 如果 axis = 0, out[i][j][k] = input[i][indices[i][j][k]][k] 如果 axis = 1, out[i][j][k] = input[i][j][indices[i][j][k]] 如果 axis = 2。

例如：

gather([[4, 2], [5, 3]], 0, [[1, 0], [0, 1]])

将生成 [[5, 2], [4, 3]]。

gather([[1, 2, 3], [4, 5, 6], 1, [[1], [0]])

将生成 [[2], [4]]。更多示例请参见 ONNX 文档。

该操作映射到 TensorRT IGatherLayer。

参数:

input – Tensor 要从中收集元素的输入张量。
dim – int 要收集的维度。
indices – Union[Tensor, int] 从“dim”维度收集的位置。

返回值:

包含收集到的元素的张量。它与索引张量具有相同的形状。

tensorrt_llm.functional.gather_last_token_logits( hidden_states: Tensor, last_token_ids: Tensor, remove_input_padding: bool, ) → Tensor[source]#

从隐藏状态中提取与最后一个 token 对应的 logits。

该函数添加操作以提取一批序列中最后一个 token 的 logits。

根据 'remove_input_padding' 是 'True' 还是 'False'，该函数假定不同形状的输入。

当 'remove_input_padding' 为 'True' 时，假定 'hidden_states' 张量是 packed 的。它的形状为 '[num_tokens, hidden_dim]'，其中 'num_tokens' 是批次中序列长度的总和，'hidden_dim' 是隐藏维度。 'last_tokens_ids' 是一个一维张量，它编码了批次中序列长度的包含前缀和。

当 'remove_input_padding' 为 'False' 时，假定 'hidden_states' 张量是 padded 的。它的形状为 '[batch_size, max_seqlen, hidden_dim]'，其中 'max_seqlen' 是批次中最长序列的长度，'hidden_dim' 是隐藏维度。 'last_token_ids' 是一个一维张量，它编码了批次中每个序列的长度。

在这两种情况下，该函数都会生成一个形状为 '[batch_size, hidden_size]' 的张量，其中索引 'i' 处的行对应于来自第 'i' 个序列的最后一个 token 的 logits。

参数:

hidden_states – Tensor 隐藏状态
last_token_ids – Tensor 批次中序列的长度或包含前缀和的长度。
remove_input_padding – bool 指示 hidden_states 是 packed ('True') 还是 padded ('False')。

返回值:

由该系列操作创建的张量。

tensorrt_llm.functional.gather_nd( input: Tensor, indices: Tensor, batch_dims: int = 1, ) → Tensor[source]#

添加一个执行 gather 操作的层，其中包含一些 element-wise 维度。请参见：https://onnx.org.cn/onnx/operators/onnx__GatherND.html gather 操作在 dim=batch_dims 上执行。

参数:

input – Tensor 执行 gather 操作的张量。
indices – Tensor 指示要收集哪些条目的张量。
batch_dims – int 在 gather 开始之前应跳过的第一个维度的数量。

返回值:

由具有 GatherMode.ND 的 gather 层创建的张量。

tensorrt_llm.functional.gegelu( x: Tensor, limit: float | None = None, ) → Tensor[source]#

tensorrt_llm.functional.geglu( x: Tensor, ) → Tensor[source]#

添加一个 Gated-GELU 操作。

该函数接受一个张量，将其沿最后一个维度分成两半，将 GELU 应用于后半部分，然后将结果相乘。如果最后一个维度不是偶数，则该行为未定义。

参数:: input – Tensor 应用激活函数的输入张量。
返回值:: 激活层生成的张量。

tensorrt_llm.functional.gelu( x: Tensor, ) → Tensor[source]#

添加一个 GELU 操作。

参数:: input – Tensor 应用激活函数的输入张量。
返回值:: 激活层生成的张量。

tensorrt_llm.functional.gemm_allreduce( a: Tensor, b: Tensor, group: List[int], transa: bool = False, transb: bool = False, alpha: ndarray | Tensor | None = None, output_dtype: DataType | None = None, fp8_inputs_override: bool = False, a_sf: Tensor | None = None, b_sf: Tensor | None = None, )[source]#

添加一个执行融合GEMM+AllReduce的操作。

参数:

a – Tensor 输入张量 A
b – Tensor 输入张量 B
a_sf – Optional[Tensor] 用于缩放输入 A 的可选输入张量
b_sf – Optional[Tensor] 用于缩放输入 B 的可选输入张量
group – List[int] 参与集合运算的 ranks 列表
transa – bool 是否转置输入张量 A
transb – bool 是否转置输入张量 B
alpha – float GEMM 的 Alpha 值 -> beta * C + (alpha * acc)
output_dtype – trt.DataType 插件的输出类型。如果为 None，我们将使用 plugin_config 中设置的类型。
fp8_inputs_override – bool TRT 图无法正确检测 FP8 输入。此标志用于覆盖导出的输入张量类型，以便我们的插件知道发出 FP8 MMA。

返回值:

返回 GEMM 输出张量，该张量已在 ranks 之间减少。

tensorrt_llm.functional.gemm_swiglu( input: Tensor, weight: Tensor, bias: Tensor | None = None, scale_d0: float = 1.0, scale_d1: float = 1.0, scale_output: float = 1.0, ) → Tensor[source]#

添加一个矩阵乘法，后跟 SwiGLU (x * SiLU(gate)) 操作。

第二个 SwiGLU 操作获取前面的张量，沿着最后一个维度将其分成两半，将 SiLU 应用于第二半，然后将结果相乘。如果最后一个维度不是偶数，则行为未定义。

参数：input : Tensor

第一个张量（通常称为 A）。

weightTensor
第二个张量（通常称为 B）。

biasOptional[Tensor]
逐通道偏差。具有 fp8 dtype 的插件尚不支持偏差。

scale_d0float
用于解量化 x 的比例，用于 fp8

scale_d1float
用于解量化 gate 的比例，用于 fp8

scale_outputfloat
用于量化输出的比例，用于 fp8

返回值

插入层生成的张量。

tensorrt_llm.functional.generate_alibi_biases( slopes: Tensor, key_length: Tensor, ) → Tensor[source]#

计算 ALiBi 偏差，如 https://arxiv.org/abs/2211.05100 中所述。

ALiBi 偏差被添加到多头注意力模块中的 Q*K^T 乘积的结果中。

参数:

slopes – Tensor 斜率。
key_length – Tensor 每个头的 K 向量的大小。

返回值:

一个包含 ALiBi 偏差的常数张量。

tensorrt_llm.functional.generate_alibi_slopes( num_heads: int, tp_size: int = 1, tp_rank: int = 0, alibi_scale: float = 1.0, alibi_bias_max: int = 8, ) → ndarray[source]#

计算 ALiBi 斜率，如 https://arxiv.org/abs/2211.05100 中所述。

参数:

num_heads – int 头数。
dtype – trt.DataType 返回的斜率的数据类型
tp_size – int 张量并行大小
tp_rank – int 张量并行 rank

返回值:

一个包含 ALiBi 斜率的常数张量。

tensorrt_llm.functional.generate_logn_scaling( seq_length: int = 8192, max_position_embeddings: int = 32768, ) → ndarray[source]#

计算 Qwen 推理外推的 Log-N 缩放向量

参数:

seq_length – int 训练中的最大序列长度 (Qwen-1 中默认为 8192)
max_position_embeddings – int 最大位置嵌入。(Qwen-1 中默认为 32768)

返回值:

一个包含 logn 缩放向量的常数 np.ndarray

tensorrt_llm.functional.gpt_attention(*, qkv: ~tensorrt_llm.functional.Tensor, past_key_value: ~tensorrt_llm.functional.Tensor, attention_mask: ~tensorrt_llm.functional.Tensor | None = None, attention_packed_mask: ~tensorrt_llm.functional.Tensor | None = None, sequence_length: ~tensorrt_llm.functional.Tensor, host_past_key_value_lengths: ~tensorrt_llm.functional.Tensor | None, host_max_attention_window_sizes: ~tensorrt_llm.functional.Tensor, host_sink_token_length: ~tensorrt_llm.functional.Tensor, context_lengths: ~tensorrt_llm.functional.Tensor | None, cache_indirection: ~tensorrt_llm.functional.Tensor | None, host_request_types: ~tensorrt_llm.functional.Tensor, layer_idx: int, num_heads: int, num_kv_heads: int, hidden_size_per_head: int, q_scaling: float, attn_logit_softcapping_scale: float = 0.0, rotary_embedding_dim: int = 0, rotary_embedding_base: float = 10000.0, rotary_embedding_scale_type: ~tensorrt_llm.functional.RotaryScalingType = RotaryScalingType.none, rotary_embedding_short_m_scale: float = 1.0, rotary_embedding_long_m_scale: float = 1.0, rotary_embedding_scale: float = 1.0, rotary_embedding_max_positions: int = 1024, rotary_embedding_original_max_positions: int = 1024, position_embedding_type: ~tensorrt_llm.functional.PositionEmbeddingType = PositionEmbeddingType.learned_absolute, rotary_inv_freq: ~tensorrt_llm.functional.Tensor | None = None, rotary_cos_sin: ~tensorrt_llm.functional.Tensor | None = None, kv_orig_quant_scale: ~tensorrt_llm.functional.Tensor | None = None, kv_quant_orig_scale: ~tensorrt_llm.functional.Tensor | None = None, attention_output_orig_quant_scale: ~tensorrt_llm.functional.Tensor | None = None, attention_output_sf_scale: ~tensorrt_llm.functional.Tensor | None = None, kv_cache_quant_mode: ~tensorrt_llm._utils.QuantModeWrapper | ~tensorrt_llm.quantization.mode.QuantMode = <QuantMode: 0>, max_context_length: int | None = None, mask_type: ~tensorrt_llm.functional.AttentionMaskType = AttentionMaskType.causal, block_sparse_block_size: int = 64, block_sparse_homo_head_pattern: bool = False, block_sparse_num_local_blocks: int = 16, block_sparse_vertical_stride: int = 8, alibi_slopes: ~tensorrt_llm.functional.Tensor | None = None, tp_size: int = 1, tp_rank: int = 0, vision_start: int = -1, vision_length: int = -1, kv_cache_block_offsets: ~tensorrt_llm.functional.Tensor | None = None, host_kv_cache_block_offsets: ~tensorrt_llm.functional.Tensor = None, host_kv_cache_pool_pointers: ~tensorrt_llm.functional.Tensor = None, host_kv_cache_pool_mapping: ~tensorrt_llm.functional.Tensor = None, do_cross_attention: bool = False, cross_kv: ~tensorrt_llm.functional.Tensor | None = None, cross_kv_length: ~tensorrt_llm.functional.Tensor | None = None, encoder_input_lengths: ~tensorrt_llm.functional.Tensor | None = None, relative_attention_bias: ~tensorrt_llm.functional.Tensor | None = None, logn_scaling: ~tensorrt_llm.functional.Tensor | None = None, max_distance: int = 0, host_context_lengths: ~tensorrt_llm.functional.Tensor | None = None, qkv_bias: ~tensorrt_llm.functional.Tensor | None = None, use_cache: bool = True, spec_decoding_is_generation_length_variable: bool = False, spec_decoding_max_generation_length: int = 0, spec_decoding_generation_lengths: ~tensorrt_llm.functional.Tensor = None, spec_decoding_position_offsets: ~tensorrt_llm.functional.Tensor = None, spec_decoding_packed_mask: ~tensorrt_llm.functional.Tensor = None, spec_decoding_use: ~tensorrt_llm.functional.Tensor = None, long_rope_rotary_inv_freq: ~tensorrt_llm.functional.Tensor | None = None, long_rope_rotary_cos_sin: ~tensorrt_llm.functional.Tensor | None = None, mrope_rotary_cos_sin: ~tensorrt_llm.functional.Tensor = None, mrope_position_deltas: ~tensorrt_llm.functional.Tensor = None, host_runtime_perf_knobs: ~tensorrt_llm.functional.Tensor | None = None, host_context_progress: ~tensorrt_llm.functional.Tensor = None, is_mla_enabled_flag: bool = False, q_lora_rank: int = 0, kv_lora_rank: int = 0, qk_nope_head_dim: int = 0, qk_rope_head_dim: int = 0, v_head_dim: int = 0, q_b_proj: ~tensorrt_llm.functional.Tensor | None = None, kv_b_proj: ~tensorrt_llm.functional.Tensor | None = None, k_b_proj_trans: ~tensorrt_llm.functional.Tensor | None = None, skip_attn=None, cp_group: ~typing.List[int] = [0], cp_size: int = 1, cp_rank: int = 0, num_kv_heads_origin: int = -1) → Tuple[Tensor, Tensor | None][source]#

添加一个在类似GPT的模型中执行多头注意力机制的操作。

该函数的签名将在未来的版本中更改 - 我们正在简化API。当前版本仍在开发中！以下API提供了一些关于未来版本中可能被删除或与其他参数合并的参数的提示。

有关该函数的文档，请参见docs/source/advanced/gpt-attention.md。

参数:

qkv – Tensor (在 GPU 上) 输入的 QKV 张量。在填充模式下，其形状为 [batch_beam_size, max_seqlen, qkv_dim]，在打包模式下为 [num_tokens, qkv_dim]。其中 qkv_dim 取决于使用 MQA、GQA 或 MHA。请参阅 docs/source/advanced/gpt-attention.md 中的 QKV 输入。
past_key_value – Tensor (在 GPU 上) 存储 KV 缓存数据的张量。在连续模式下，其形状为 [max_batch_size * max_beam_width, 2, num_kv_heads, max_seqlen, hidden_dim_per_head]，在分页模式下为 [max_blocks, 2, num_kv_heads, num_tokens_per_block, hidden_dim_per_head]。请参阅 docs/source/advanced/gpt-attention.md 中的 KV 缓存。
attention_mask – Tensor (在 GPU 上) 存储用于非融合 MHA 或 MMHA 的注意力掩码的张量。其形状为 [num_tokens, max_kv_seqlen]。
attention_packed_mask – Tensor (在 GPU 上) 存储用于 fmha 的打包自定义掩码的张量。其形状为 [num_tokens, max_kv_seqlen / 32]，其中每个位代表一个掩码位置。
sequence_lengths – Tensor (位于 GPU 上) 用于存储每个序列长度的张量。其形状为 [batch_size]。参见 docs/source/advanced/gpt-attention.md 中的 QKV 输入。
host_past_key_value_lengths – Tensor (位于 CPU 上) 形状为 [batch_size] 的 INT32 张量。
host_max_attention_window_sizes – Tensor (位于 CPU 上) 形状为 [1] 的 INT32 张量。默认情况下，max_attention_window_size 由 cache_indir_table 的形状决定。我们支持每个层独立设置 max_attention_window_size。这控制了滑动窗口注意力/循环 KV 缓存特性。
context_lengths – Tensor (位于 GPU 上) 用于存储每个请求的上下文阶段序列长度的张量。其形状为 [batch_size]。参见 doc/functional.py 中的 QKV 输入。
cache_indirection – Tensor (位于 GPU 上) 用于在使用集束搜索时重建路径的张量。其形状为 [batch_size, beam_width, max_seqlen]。参见 docs/source/advanced/gpt-attention.md 中的集束搜索。
host_request_types – Tensor = None (位于 CPU 上) 位于主机上的张量，指示请求是否处于上下文或生成阶段。其形状为 [batch_size]。参见 docs/source/advanced/gpt-attention.md 中的 Inflight Batching。
layer_idx – int 此注意力层的索引，用于访问 kv_cache_block_offsets。
num_heads – int 头的数量。
num_kv_heads – int KV 头的数量，通用处理 MHA/MQA/GQA。
hidden_size_per_head – int 每个头的隐藏大小。
q_scaling – float 用于计算应用于 Q*K^T 乘积输出的缩放因子的值。参见 docs/source/advanced/gpt-attention.md 中的缩放因子。
attn_logit_softcapping_scale – float 用于计算应用于 Q*K^T 乘积输出的缩放因子的 scale * tanh(value / scale)。
rotary_embedding_dim – int 用于计算 RoPE 的维度。当 position_embedding_type 不是 RoPE 时，使用 0。
rotary_embedding_base – float 用于 RoPE 的 theta 值。当 position_embedding_type 不是 RoPE 时，忽略。
rotary_embedding_scale_type –
RotaryScalingType RoPE 的缩放类型。当 position_embedding_type 不是 RoPE 时，忽略。可能的旋转缩放类型：
- RotaryScalingType.none
- RotaryScalingType.linear
- RotaryScalingType.dynamic
- RotaryScalingType.longrope
- RotaryScalingType.llama3
rotary_embedding_scale – float 用于 RoPE 中线性/动态缩放的缩放值。当 rotary_embedding_scale_type 为 none 时，忽略。如果 rotary_embedding_scale_type 为 none，则必须设置为 1（默认）。
rotary_inv_freq – float Tensor 形状为 [head_size / 2] 的旋转逆频率张量。
rotary_cos_sin – float2(cos/sin) Tensor 旋转 cos/sin 缓存，将在不同的请求之间重用。它被视为常量张量。
rotary_embedding_max_positions – int 仅当使用 dynamic RoPE 缩放时需要。否则忽略。
position_embedding_type –
PositionEmbeddingType 位置嵌入类型
- PositionEmbeddingType.learned_absolute
- PositionEmbeddingType.relative
- PositionEmbeddingType.rope_gptj
- PositionEmbeddingType.rope_gpt_neox
- PositionEmbeddingType.alibi
- PositionEmbeddingType.alibi_with_scale
kv_orig_quant_scale – Tensor 用于存储 KV 缓存中量化为 INT8/FP8 的缩放因子的张量。其形状为 [1]。参见 docs/source/advanced/gpt-attention.md 中的 INT8/FP8 KV 缓存。
kv_quant_orig_scale – Tensor 用于存储 KV 缓存中从 INT8/FP8 反量化的缩放因子的张量。其形状为 [1]。参见 docs/source/advanced/gpt-attention.md 中的 INT8/FP8 KV 缓存。
attention_output_orig_quant_scale – Tensor 用于存储 KV 缓存中量化为 FP8 的缩放因子的张量。其形状为 [1]。
kv_cache_quant_mode – QuantMode (int 标志) 我们是否启用 INT8 或 FP8 KV 缓存？
max_context_length – int32_t 最长输入序列的长度。参见 docs/source/advanced/gpt-attention.md 中的 QKV 输入。
mask_type –
int = 1 掩码的类型：
- tensorrt_llm.layers.AttentionMaskType.padding 用于 BERT，
- tensorrt_llm.layers.AttentionMaskType.causal 用于 GPT，
- tensorrt_llm.layers.AttentionMaskType.sliding_window_causal 用于 GPT，
- tensorrt_llm.layers.AttentionMaskType.bidirectional 用于 ChatGLM-6B，
- tensorrt_llm.layers.AttentionMaskType.bidirectionalglm 用于 GLM-10B，
- tensorrt_llm.layers.AttentionMaskType.blocksparse 用于 Phi-3-small，
- tensorrt_llm.layers.AttentionMaskType.custom_mask 用于任何模型。
block_sparse_block_size – int 块稀疏注意力中的块大小
block_sparse_homo_head_pattern – bool 所有注意力头是否共享相同的垂直步幅模式？
block_sparse_num_local_blocks – int 对角线附近活动块的数量
block_sparse_vertical_stride – int 垂直维度上活动块的步幅
alibi_slopes – Tensor ALiBi 斜率。ALiBi 偏差在可能的情况下在内核中动态计算。
tp_size – int 激活张量并行时的进程/GPU 数量。
tp_rank – int 该进程的 rank（当运行张量并行时）。
kv_cache_block_offsets – KV 缓存的块偏移张量。其形状为 [num_layers, max_batch_size, max_beam_width, 2, max_blocks_per_sequence * 2]，参见 docs/source/advanced/gpt-attention.md 中的 KV 缓存部分，位于 GPU 上。
host_kv_cache_block_offsets – 与 kv_cache_block_offsets 相同，但位于 CPU 上。
host_kv_cache_pool_pointers – KV 缓存的池指针张量。其形状为 [num_layers, 2]，参见 docs/source/advanced/gpt-attention.md 中的 KV 缓存部分，位于 GPU 上。
host_kv_cache_pool_mapping – 用于不同内存池的池映射张量。其形状为 [num_layers,2] - 对于每一层，池的索引和池内层的索引。
do_cross_attention – bool = False 我们是否使用它作为交叉注意力而不是自注意力？
cross_kv – Tensor = None 编码器输出隐藏状态的 KV 张量。其形状在填充模式下为 [batch_size, max_seqlen, 2 * kvHeadNum * headSize]，在打包模式下为 [1, num_tokens, 2 * kvHeadNum * headSize]。
cross_kv_length – Tensor = None 最长编码器输出序列的长度。
encoder_input_lengths – Tensor 用于存储每个编码器输入序列长度的张量。其形状为 [batch_size]。
logn_scaling – Tensor = None logn 缩放张量 [max_position_embedding_len]，应用于 q 以帮助外推。
relative_attention_bias – Tensor = None 相对注意力偏差 [num_heads, max_seq_len, max_seq_len]，或隐式模式的相对注意力嵌入表 [num_heads, num_buckets]。
max_distance – int = 0 注意力中相对位置的最大距离，用于隐式模式。默认值为 0，表示使用相对注意力偏差的常规模式。仅当传入非零的正 max_distance 值时，才会启用隐式模式。请参阅 docs/source/advanced/gpt-attention.md 中的相对注意力偏差
host_context_lengths – Tensor = None (位于 CPU 上) 包含不同输入长度的主机张量。
qkv_bias – Tensor = None, qkv 偏差张量。
use_cache – bool = False 我们是否需要存储 kv 缓存？如果没有生成阶段，则不需要。
spec_decoding_is_generation_length_variable – bool = False，批量处理中每个序列的生成长度是否可以不同。对于 Medusa，应设置为 False。对于 Redrafter，应设置为 True。
spec_decoding_max_generation_length – int = 1，每个序列在生成阶段可能的最大 token 数。
spec_decoding_generation_lengths – Tensor = None，每个序列的生成阶段 token 长度。形状：[batch_size]
spec_decoding_position_offsets – Tensor = None，推测解码 tokens 的位置偏移量（所有序列共享）。形状：[batch_size, num_draft_tokens + 1]。
spec_decoding_packed_mask –
Tensor = None，推测解码 tokens 的注意力掩码（打包成 uint32_t 位）。remove_input_padding 为 False

形状：[batch_size, num_draft_tokens + 1, divUp(num_draft_tokens + 1, 32)]。

remove_input_padding 为 True
形状：[sum(spec_decoding_generation_lengths), divUp(num_draft_tokens + 1, 32)]。
long_rope_rotary_inv_freq – float Tensor 用于更长序列长度的额外旋转逆频率。形状：[head_size / 2]
long_rope_rotary_cos_sin – float2(cos/sin) Tensor 用于更长序列长度的额外旋转 cos/sin 缓存。
is_mla_enable – bool = False 我们是否需要启用 deepseekv2 mla？
host_runtime_perf_knobs – Tensor = None，运行时性能旋钮位掩码，控制是否在运行时使用某些性能旋钮。
host_context_progress – Tensor = None，用于跟踪上下文阶段中逐层进度的结构。
skip_attn – Tensor = None, CPU 上的布尔张量。如果为 true，则不运行 attention 插件，直接返回。
num_kv_heads_origin – int 未经过 TP 处理的原始 KV 头数

返回值:

该层生成的张量。

tensorrt_llm.functional.group_norm( input: Tensor, num_groups: int, weight: Tensor | None = None, bias: Tensor | None = None, eps: float = 1e-05, )[source]#

tensorrt_llm.functional.gt( left: ~tensorrt_llm.functional.Tensor | int | float, right: ~tensorrt_llm.functional.Tensor | int | float, *, op: ~tensorrt.tensorrt.ElementWiseOperation = <ElementWiseOperation.GREATER: 12>, ) → Tensor#

添加具有两个输入的逐元素运算。

以下闭包在 functional.* 中定义。

add for op=trt.ElementWiseOperation.SUM sub for op=trt.ElementWiseOperation.SUB mul for op=trt.ElementWiseOperation.PROD div for op=trt.ElementWiseOperation.DIV floordiv for op=trt.ElementWiseOperation.FLOOR_DIV gt for op=trt.ElementWiseOperation.GREATER lt for op=trt.ElementWiseOperation.LESS op_and for op=trt.ElementWiseOperation.AND op_or for op=trt.ElementWiseOperation.OR eq for op=trt.ElementWiseOperation.EQUAL minimum for op=trt.ElementWiseOperation.MIN maximum for op=trt.ElementWiseOperation.MAX pow for op=trt.ElementWiseOperation.POW

它使用 TensorRT 中的 IElementWiseLayer 实现。

参数:

left – Union[Tensor, int, float] 第一个输入。如果该输入是整数或浮点数，则该函数会创建一个常量张量。
right – Union[Tensor, int, float] 第二个输入。如果该输入是整数或浮点数，则该函数会创建一个常量张量。
op – trt.ElementWiseOperation 要执行的二元运算。

返回值:

此逐元素运算生成的张量。

tensorrt_llm.functional.identity( input: Tensor, ) → Tensor[source]#

添加一个 identity 操作。

TODO: 记录为什么可以使用插件来完成!!!

参数:: input – Tensor 输入张量。
返回值:: 此 identity 操作生成的张量。

tensorrt_llm.functional.index_select( input: Tensor, dim: int, index: Tensor, ) → Tensor[source]#

添加一个从张量中选择元素切片的操作。

给定一个输入张量，该函数创建一个操作，该操作在维度“dim”处选择“index”中列出的索引处的元素切片，以创建一个新张量。输出张量与输入张量具有相同的 rank。

“index”是一个秩为 1 的张量。

例如，对于输入=[[4, 2, 5], [2, 1, 2], [4, 7, 1]]，其形状为 [3, 3]，

index_select(input, 0, [0, 1])

将创建一个形状为 [2, 3] 的张量，其中包含 [[4, 2, 5], [2, 1, 2]]。

关于输出张量的形状，维度 ‘dim’ 与 ‘index’ 张量具有相同的大小。这意味着对于形状为 [4, 2, 6, 3] 的输入张量，

index_select(input, 2, [1, 4])

将从第三个维度（dim == 2）中选择第 2 个和第 5 个切片（index == 1 或 4），并返回形状为 [4, 2, 2, 3] 的张量（即，第三个维度缩小到 2）。

请注意，此操作也可用于在 ‘dim’ 维度中扩展张量，例如，在输入 [[0, 1], [2, 3]] 上，

index_select(input, 1, [0, 0, 0])

将生成一个形状为 [2, 3] 的张量，其中包含 [[0, 0, 0], [2, 2, 2]]。

该操作映射到 TensorRT IGatherLayer。

参数:

input – Tensor 要从中选择的输入张量。
dim – int 要从中选择的维度。
index – Tensor 要选择的 ‘dim’ 维度中切片的索引。

返回值:

包含所选切片的张量。

tensorrt_llm.functional.int_clip( input: Tensor, lower: int, upper: int, ) → Tensor[source]#

tensorrt_llm.functional.interpolate( input: Tensor, size: int | List[int] = None, scale_factor: float | List[float] = None, mode: str = 'nearest', align_corners: bool = False, recompute_scale_factor: bool = False, antialias: bool = False, ) → Tensor[source]#

tensorrt_llm.functional.is_gated_activation(activation)[source]#

给定的激活函数是否为门控激活函数？

参数:: activation – str 激活函数的名称。
返回值:: 如果函数是门控激活函数，则为 True，否则为 False。

tensorrt_llm.functional.layer_norm( input: Tensor, normalized_shape: int | Tuple[int], weight: Tensor | None = None, bias: Tensor | None = None, eps: float = 1e-05, use_diff_of_squares: bool = True, ) → Tensor[source]#

在张量上添加一个层归一化操作。

该操作将层归一化应用于其输入张量。在其最简单的形式中，对于大型语言模型，“normalized_shape”应设置为激活张量的隐藏维度。否则，它是张量的归一化部分（从最右边的维度开始）的形状。

“weight”张量对应于层归一化公式中的“gamma”，“bias”是“beta”。在计算平方根之前，将“eps”值添加到方差中。

此实现（使用插件时）支持一个额外的标志，用于启用/禁用平方差的使用（“Var = Mean(X^2) - Mean(X)^2”）。

参数:

input – Tensor 要归一化的张量。
normalized_shape – Union[int, Tuple[int]] 归一化的子张量的形状。使用“hidden_dim”来归一化 LLM 中激活张量的最内层维度。
weight – Optional[Tensor] = None 层归一化中的 “gamma” 项。它的形状必须是 ‘normalized_shape’。
bias – Optional[Tensor] = None 层归一化中的 “beta” 项。它的形状必须是 ‘normalized_shape’。
eps – float 要添加到平方根方差中的 epsilon 项。
use_diff_of_squares – bool 插件是否使用平方差来计算方差？

返回值:

该操作的输出张量。

tensorrt_llm.functional.log( input: ~tensorrt_llm.functional.Tensor, *, op: ~tensorrt.tensorrt.UnaryOperation = <UnaryOperation.LOG: 1>, ) → Tensor#

对单个输入添加逐元素运算。

以下闭包在 functional.* 中定义。

round for op=trt.UnaryOperation.ROUND sqrt for op=trt.UnaryOperation.SQRT exp for op=trt.UnaryOperation.EXP sin for op=trt.UnaryOperation.SIN cos for op=trt.UnaryOperation.COS abs for op=trt.UnaryOperation.ABS log for op=trt.UnaryOperation.LOG

它使用 TensorRT 中的 IUnaryLayer 实现。

参数:

input – Tensor 输入张量。
op – trt.UnaryOperation 要执行的一元运算。

返回值:

此逐元素运算生成的张量。

tensorrt_llm.functional.log_softmax( input: Tensor, dim: int, ) → Tensor[source]#

此函数等效于 torch.nn.functional.log_softmax()，即它以更安全、更快速的方式执行 log(softmax(input))。

参数:

input – Tensor 要在其上计算 log_softmax 的数据张量。
dim – int 输入张量的维度，将沿该维度计算 log_softmax。

返回值:

与输入形状相同的张量，并且在指定的 dim 上计算 log_softmax。

tensorrt_llm.functional.lora_plugin( input: Tensor = None, in_hidden_size: int = 0, out_hidden_sizes: List[int] = [0], host_request_types: Tensor = None, transa: bool = False, transb: bool = False, host_context_lengths: Tensor = None, max_low_rank: int = 0, lora_ranks: List[Tensor] = None, lora_weights_pointers: List[Tensor] = None, weight_index: int = 0, )[source]#

参数:

input – Tensor (在 GPU 上) 输入张量。其形状为 [batch_size, seq_len, dim] 或 [num_tokens, dim]（用于 remove_input_padding）
in_hidden_size/out_hidden_size – int lora 计算工作流程是 [M, in_hidden_size] -> [M, low_rank] -> [M, out_hidden_size]
host_request_types – Tensor = None 主机上的张量，指示请求处于上下文阶段还是生成阶段。其形状为 [batch_size]。请参阅 docs/source/advanced/gpt-attention.md 中的 Inflight Batching。
transa – bool 第一个输入是否转置？如果要转置第一个输入，则设置为“True”，否则设置为“False”。
transb – bool 是否转置第二个输入？如果要转置第二个输入，则设置为“True”，否则设置为“False”。
host_context_lengths – cpu Tensor = None 包含不同输入长度的主机张量。
max_low_rank – int 最大 low_rank，用于确定工作区大小。
lora_ranks – 形状为 [batch_size] 的 cpu Tensor 每个请求的 low_rank
lora_weights_pointers – 形状为 [batch_size, 3] 的 cpu int64 Tensor。每个请求的权重指针。由 in_pointer、out_pointer 以及可能的 scales 向量指针组成。
weight_index – int 如果权重指针指向多个权重，则为权重的索引。

返回值:

该层生成的张量。

tensorrt_llm.functional.low_latency_gemm( input: Tensor, mat2: Tensor, alpha: ndarray | None = None, strict_dtype: DataType | None = None, ) → Tensor[source]#

tensorrt_llm.functional.low_latency_gemm_swiglu( input: Tensor, weight: Tensor, scale_d0: float = 1.0, scale_d1: float = 1.0, scale_output: float = 1.0, ) → Tensor[source]#

添加一个矩阵乘法，后跟 SwiGLU (x * SiLU(gate)) 操作。

参数：input : Tensor

第一个张量（通常称为 A）。

weightTensor
第二个张量（通常称为 B）。

scale_d0float
用于解量化 x 的比例，用于 fp8

scale_d1float
用于解量化 gate 的比例，用于 fp8

scale_outputfloat
用于量化输出的比例，用于 fp8

返回值

插入层生成的张量。

tensorrt_llm.functional.lt( left: ~tensorrt_llm.functional.Tensor | int | float, right: ~tensorrt_llm.functional.Tensor | int | float, *, op: ~tensorrt.tensorrt.ElementWiseOperation = <ElementWiseOperation.LESS: 13>, ) → Tensor#

添加具有两个输入的逐元素运算。

以下闭包在 functional.* 中定义。

add for op=trt.ElementWiseOperation.SUM sub for op=trt.ElementWiseOperation.SUB mul for op=trt.ElementWiseOperation.PROD div for op=trt.ElementWiseOperation.DIV floordiv for op=trt.ElementWiseOperation.FLOOR_DIV gt for op=trt.ElementWiseOperation.GREATER lt for op=trt.ElementWiseOperation.LESS op_and for op=trt.ElementWiseOperation.AND op_or for op=trt.ElementWiseOperation.OR eq for op=trt.ElementWiseOperation.EQUAL minimum for op=trt.ElementWiseOperation.MIN maximum for op=trt.ElementWiseOperation.MAX pow for op=trt.ElementWiseOperation.POW

它使用 TensorRT 中的 IElementWiseLayer 实现。

参数:

left – Union[Tensor, int, float] 第一个输入。如果该输入是整数或浮点数，则该函数会创建一个常量张量。
right – Union[Tensor, int, float] 第二个输入。如果该输入是整数或浮点数，则该函数会创建一个常量张量。
op – trt.ElementWiseOperation 要执行的二元运算。

返回值:

此逐元素运算生成的张量。

tensorrt_llm.functional.mamba_conv1d( input: Tensor, conv_state_or_ptr: Tensor, conv_weight: Tensor, conv_bias: Tensor, host_request_types: Tensor, last_token_ids: Tensor, dim: int, dconv: int, dtype: str, pre_stride: int = 0, post_stride: int = 0, host_context_lengths: Tensor | None = None, slot_mapping: Tensor | None = None, apply_silu: bool = True, )[source]#

参数:

input – Tensor (在 GPU 上) 输入张量。其形状为 [batch_size, seq_len, dim] 或 [num_tokens, dim]（用于 remove_input_padding）
conv_state_or_ptr – Tensor（在 GPU 或 CPU 上）conv 状态张量。其形状为 [batch_size, dconv - 1, dim] 或形状为 [1] 的 CPU 张量，用于指向分页状态。
conv_weight – Tensor（在 GPU 上）权重张量。其形状为 [1, dconv, dim]
conv_bias – Tensor（在 GPU 上）偏置张量。其形状为 [dim]
host_request_types – Tensor（在 CPU 上）主机上的张量，指示请求是在上下文还是生成阶段。其形状为 [batch_size]。请参阅 docs/source/advanced/gpt-attention.md 中的 Inflight Batching。
last_token_ids – Tensor（在 GPU 上）批次中序列长度的包含前缀和或序列的长度。
dim – int conv1d 的隐藏维度
dconv – int conv1d 的窗口大小
dtype – str 数据类型
pre_stride – int = 0 输入张量的（前）步幅大小。输入张量的有效值为 input[…, pre_stride: dim-post_stride]
post_stride – int = 0 输入张量的（后）步幅大小。输入张量的有效值为 input[…, pre_stride: dim-post_stride]
host_context_lengths – Tensor (在 CPU 上) (可选) 包含不同输入长度的主机张量。
slot_mapping – Tensor (在 GPU 上) (可选) 状态中的实际页面索引。其形状为 [dim]，用于分页状态，每页形状为 [dconv, dim]
apply_silu – bool conv1d 之后是否有 SiLU 操作？如果为 True，则在 conv1d 之后应用 SiLU 激活函数。

tensorrt_llm.functional.masked_scatter( input: Tensor, mask: Tensor, source: Tensor, ) → Tensor[source]#

基于 PyTorch 定义添加 masked_scatter。

请参阅 https://pytorch.ac.cn/docs/stable/generated/torch.Tensor.masked_scatter_.html#torch-tensor-masked-scatter 以获取该函数的描述。

参数:

input – Tensor 输入张量。
mask – Tensor 指示要选择的元素的布尔掩码张量。
source – Tensor 要从中复制的张量

返回值:

包含由掩码选择的源张量的张量。

tensorrt_llm.functional.masked_select( input: Tensor, mask: Tensor, ) → Tensor[source]#

添加一个操作，用于根据布尔掩码张量从张量中选择元素。

给定一个输入张量，该函数创建一个操作，该操作选择由布尔掩码张量指示的索引处的元素，以创建一个新的张量。输出张量是一个 1-D 张量。

输入张量必须具有 rank >= 1。输入张量和掩码张量的形状不需要匹配，但它们必须能够广播。

例如，对于输入=[[4, 2, 5], [2, 1, 2], [4, 7, 1]]，其形状为 [3, 3]，

masked_select(input, [[True, False, True], [False, True, False], [True, False, True]])

将创建一个形状为 [5] 的张量，其中包含 [4, 5, 1, 4, 1]。

masked_select(input, [[True], [False], [True]])

将创建一个形状为 [6] 的张量，其中包含 [4, 2, 5, 4, 7, 1]。

masked_select(input, [[False, False, True]])

将创建一个形状为 [3] 的张量，其中包含 [5, 2, 1]。

masked_select(input, [False])

将创建一个形状为 [0] 的张量，该张量为空。

该操作由 TensorRT 中的 NonZero、Shuffle 和 GatherV2 层实现。

参数:

input – Tensor 要从中选择的输入张量。
mask – Tensor 指示要选择的元素的布尔掩码张量。

返回值:

包含所选元素的 1-D 张量。

tensorrt_llm.functional.matmul( input: Tensor, mat2: Tensor, transa: bool = False, transb: bool = False, use_fp32_acc: bool = True, ) → Tensor[source]#

添加矩阵乘法。

该操作映射到 tensorrt.IMatrixMultiplyLayer 层。正如 TensorRT 文档中所解释的那样，它在对输入应用可选的转置后计算两个输入之间的内积。

参数:

input – Tensor 第一个张量（通常称为 A）。
mat2 – Tensor 第二个张量（通常称为 B）。
transa – bool 第一个输入是否转置？如果要转置第一个输入，则设置为“True”，否则设置为“False”。
transb – bool 是否转置第二个输入？如果要转置第二个输入，则设置为“True”，否则设置为“False”。
use_fp32_acc – bool 如果出于准确性原因，此 fp16 matmul 需要使用 fp32 累积，则设置为“True”。这可以是每个模型和每个 matmul 的决定。

返回值:

插入层生成的张量。

tensorrt_llm.functional.max( input: Tensor, dim: int, keepdim: bool = False, ) → Tensor[source]#

添加一个操作以计算沿维度的最大值。

计算输入张量的维度“dim”上的最大值。

它使用 TensorRT 中的 IReduceLayer 实现。

参数:

input – Tensor 输入张量。
dim – int 计算平均值的维度。
keepdim – bool 该维度是否保留在缩减的张量中？如果为 True，则保留该维度，否则将其从形状中删除。

返回值:

此缩减操作生成的张量。

tensorrt_llm.functional.maximum( left: ~tensorrt_llm.functional.Tensor | int | float, right: ~tensorrt_llm.functional.Tensor | int | float, *, op: ~tensorrt.tensorrt.ElementWiseOperation = <ElementWiseOperation.MAX: 2>, ) → Tensor#

添加具有两个输入的逐元素运算。

以下闭包在 functional.* 中定义。

add for op=trt.ElementWiseOperation.SUM sub for op=trt.ElementWiseOperation.SUB mul for op=trt.ElementWiseOperation.PROD div for op=trt.ElementWiseOperation.DIV floordiv for op=trt.ElementWiseOperation.FLOOR_DIV gt for op=trt.ElementWiseOperation.GREATER lt for op=trt.ElementWiseOperation.LESS op_and for op=trt.ElementWiseOperation.AND op_or for op=trt.ElementWiseOperation.OR eq for op=trt.ElementWiseOperation.EQUAL minimum for op=trt.ElementWiseOperation.MIN maximum for op=trt.ElementWiseOperation.MAX pow for op=trt.ElementWiseOperation.POW

它使用 TensorRT 中的 IElementWiseLayer 实现。

参数:

left – Union[Tensor, int, float] 第一个输入。如果该输入是整数或浮点数，则该函数会创建一个常量张量。
right – Union[Tensor, int, float] 第二个输入。如果该输入是整数或浮点数，则该函数会创建一个常量张量。
op – trt.ElementWiseOperation 要执行的二元运算。

返回值:

此逐元素运算生成的张量。

tensorrt_llm.functional.mean( input: Tensor, dim: int | Tuple[int], keepdim: bool = False, ) → Tensor[source]#

添加一个沿着维度计算均值的操作。

计算输入张量沿维度 “dim” 的均值。

它使用 TensorRT 中的 IReduceLayer 实现。

参数:

input – Tensor 输入张量。
dim – int 计算平均值的维度。
keepdim – bool 该维度是否保留在缩减的张量中？如果为 True，则保留该维度，否则将其从形状中删除。

返回值:

此缩减操作生成的张量。

tensorrt_llm.functional.meshgrid2d( x: Tensor, y: Tensor, ) → Tuple[Tensor][source]#

创建由一维输入指定的坐标网格（二维）（仅支持indexing='xy'）。

参数:

x – Tensor 第一个输入（一维）张量。
y – Tensor 第二个输入（一维）张量。

返回值:

产生的两个张量的元组。

待办事项：添加对 torch.meshgrid 的完整支持。: 参见 https://pytorch.ac.cn/docs/stable/generated/torch.meshgrid.html#torch-meshgrid

tensorrt_llm.functional.min(input: ~tensorrt_llm.functional.Tensor, *, op: ~tensorrt.tensorrt.ReduceOperation = <ReduceOperation.MIN: 3>, dim: int | ~typing.Tuple[int], keepdim: bool = False) → Tensor#

添加一个沿维度进行的归约操作。

它使用 TensorRT 中的 IReduceLayer 实现。

参数:

input – Tensor 输入张量。
op – trt.ReduceOperation 要执行的归约操作。选项：SUM、PROD、MAX、MIN、AVG
dim – int 执行归约的维度。
keepdim – bool 该维度是否保留在缩减的张量中？如果为 True，则保留该维度，否则将其从形状中删除。

返回值:

此缩减操作生成的张量。

tensorrt_llm.functional.minimum( left: ~tensorrt_llm.functional.Tensor | int | float, right: ~tensorrt_llm.functional.Tensor | int | float, *, op: ~tensorrt.tensorrt.ElementWiseOperation = <ElementWiseOperation.MIN: 3>, ) → Tensor#

添加具有两个输入的逐元素运算。

以下闭包在 functional.* 中定义。

add for op=trt.ElementWiseOperation.SUM sub for op=trt.ElementWiseOperation.SUB mul for op=trt.ElementWiseOperation.PROD div for op=trt.ElementWiseOperation.DIV floordiv for op=trt.ElementWiseOperation.FLOOR_DIV gt for op=trt.ElementWiseOperation.GREATER lt for op=trt.ElementWiseOperation.LESS op_and for op=trt.ElementWiseOperation.AND op_or for op=trt.ElementWiseOperation.OR eq for op=trt.ElementWiseOperation.EQUAL minimum for op=trt.ElementWiseOperation.MIN maximum for op=trt.ElementWiseOperation.MAX pow for op=trt.ElementWiseOperation.POW

它使用 TensorRT 中的 IElementWiseLayer 实现。

参数:

left – Union[Tensor, int, float] 第一个输入。如果该输入是整数或浮点数，则该函数会创建一个常量张量。
right – Union[Tensor, int, float] 第二个输入。如果该输入是整数或浮点数，则该函数会创建一个常量张量。
op – trt.ElementWiseOperation 要执行的二元运算。

返回值:

此逐元素运算生成的张量。

tensorrt_llm.functional.modulo( x: Tensor, y: Tensor | int, ) → Tensor[source]#

此函数为给定的张量添加一个逐元素的模运算（x % y）。由于没有可以直接执行此操作的 TensorRT 层，因此此函数使用一些基本操作来实现它。

返回值:: 表示 (x % y) 模运算的张量。

tensorrt_llm.functional.mul( left: ~tensorrt_llm.functional.Tensor | int | float, right: ~tensorrt_llm.functional.Tensor | int | float, *, op: ~tensorrt.tensorrt.ElementWiseOperation = <ElementWiseOperation.PROD: 1>, ) → Tensor#

添加具有两个输入的逐元素运算。

以下闭包在 functional.* 中定义。

add for op=trt.ElementWiseOperation.SUM sub for op=trt.ElementWiseOperation.SUB mul for op=trt.ElementWiseOperation.PROD div for op=trt.ElementWiseOperation.DIV floordiv for op=trt.ElementWiseOperation.FLOOR_DIV gt for op=trt.ElementWiseOperation.GREATER lt for op=trt.ElementWiseOperation.LESS op_and for op=trt.ElementWiseOperation.AND op_or for op=trt.ElementWiseOperation.OR eq for op=trt.ElementWiseOperation.EQUAL minimum for op=trt.ElementWiseOperation.MIN maximum for op=trt.ElementWiseOperation.MAX pow for op=trt.ElementWiseOperation.POW

它使用 TensorRT 中的 IElementWiseLayer 实现。

参数:

left – Union[Tensor, int, float] 第一个输入。如果该输入是整数或浮点数，则该函数会创建一个常量张量。
right – Union[Tensor, int, float] 第二个输入。如果该输入是整数或浮点数，则该函数会创建一个常量张量。
op – trt.ElementWiseOperation 要执行的二元运算。

返回值:

此逐元素运算生成的张量。

tensorrt_llm.functional.non_gated_version(activation)[source]#

给定一个激活函数，获取非门控版本。

如果激活函数是非门控的，它会返回相同的激活函数名称。

例如，该函数为“swiglu”返回“silu”，为“relu”返回“relu”。

参数:: activation – str 激活函数的名称。
返回值:: 非门控激活函数的名称。

tensorrt_llm.functional.nonzero( input: Tensor, ) → Tensor[source]#

添加一个查找输入张量的非零值索引的层。

参数:: input – Tensor 我们需要查找非零值索引的输入张量。
返回值:: 形状为 [D, C] 的张量，其中 D 是 input 的维度数，C 是其中的非零值数。此 2D 张量的每一列代表每个非零值的索引元组。

tensorrt_llm.functional.not_op( input: ~tensorrt_llm.functional.Tensor, *, op: ~tensorrt.tensorrt.UnaryOperation = <UnaryOperation.NOT: 20>, ) → Tensor#

对单个输入添加逐元素运算。

以下闭包在 functional.* 中定义。

round for op=trt.UnaryOperation.ROUND sqrt for op=trt.UnaryOperation.SQRT exp for op=trt.UnaryOperation.EXP sin for op=trt.UnaryOperation.SIN cos for op=trt.UnaryOperation.COS abs for op=trt.UnaryOperation.ABS log for op=trt.UnaryOperation.LOG

它使用 TensorRT 中的 IUnaryLayer 实现。

参数:

input – Tensor 输入张量。
op – trt.UnaryOperation 要执行的一元运算。

返回值:

此逐元素运算生成的张量。

tensorrt_llm.functional.op_and( left: ~tensorrt_llm.functional.Tensor | int | float, right: ~tensorrt_llm.functional.Tensor | int | float, *, op: ~tensorrt.tensorrt.ElementWiseOperation = <ElementWiseOperation.AND: 8>, ) → Tensor#

添加具有两个输入的逐元素运算。

以下闭包在 functional.* 中定义。

add for op=trt.ElementWiseOperation.SUM sub for op=trt.ElementWiseOperation.SUB mul for op=trt.ElementWiseOperation.PROD div for op=trt.ElementWiseOperation.DIV floordiv for op=trt.ElementWiseOperation.FLOOR_DIV gt for op=trt.ElementWiseOperation.GREATER lt for op=trt.ElementWiseOperation.LESS op_and for op=trt.ElementWiseOperation.AND op_or for op=trt.ElementWiseOperation.OR eq for op=trt.ElementWiseOperation.EQUAL minimum for op=trt.ElementWiseOperation.MIN maximum for op=trt.ElementWiseOperation.MAX pow for op=trt.ElementWiseOperation.POW

它使用 TensorRT 中的 IElementWiseLayer 实现。

参数:

left – Union[Tensor, int, float] 第一个输入。如果该输入是整数或浮点数，则该函数会创建一个常量张量。
right – Union[Tensor, int, float] 第二个输入。如果该输入是整数或浮点数，则该函数会创建一个常量张量。
op – trt.ElementWiseOperation 要执行的二元运算。

返回值:

此逐元素运算生成的张量。

tensorrt_llm.functional.op_or( left: ~tensorrt_llm.functional.Tensor | int | float, right: ~tensorrt_llm.functional.Tensor | int | float, *, op: ~tensorrt.tensorrt.ElementWiseOperation = <ElementWiseOperation.OR: 9>, ) → Tensor#

添加具有两个输入的逐元素运算。

以下闭包在 functional.* 中定义。

add for op=trt.ElementWiseOperation.SUM sub for op=trt.ElementWiseOperation.SUB mul for op=trt.ElementWiseOperation.PROD div for op=trt.ElementWiseOperation.DIV floordiv for op=trt.ElementWiseOperation.FLOOR_DIV gt for op=trt.ElementWiseOperation.GREATER lt for op=trt.ElementWiseOperation.LESS op_and for op=trt.ElementWiseOperation.AND op_or for op=trt.ElementWiseOperation.OR eq for op=trt.ElementWiseOperation.EQUAL minimum for op=trt.ElementWiseOperation.MIN maximum for op=trt.ElementWiseOperation.MAX pow for op=trt.ElementWiseOperation.POW

它使用 TensorRT 中的 IElementWiseLayer 实现。

参数:

left – Union[Tensor, int, float] 第一个输入。如果该输入是整数或浮点数，则该函数会创建一个常量张量。
right – Union[Tensor, int, float] 第二个输入。如果该输入是整数或浮点数，则该函数会创建一个常量张量。
op – trt.ElementWiseOperation 要执行的二元运算。

返回值:

此逐元素运算生成的张量。

tensorrt_llm.functional.op_xor( left: ~tensorrt_llm.functional.Tensor | int | float, right: ~tensorrt_llm.functional.Tensor | int | float, *, op: ~tensorrt.tensorrt.ElementWiseOperation = <ElementWiseOperation.XOR: 10>, ) → Tensor#

添加具有两个输入的逐元素运算。

以下闭包在 functional.* 中定义。

add for op=trt.ElementWiseOperation.SUM sub for op=trt.ElementWiseOperation.SUB mul for op=trt.ElementWiseOperation.PROD div for op=trt.ElementWiseOperation.DIV floordiv for op=trt.ElementWiseOperation.FLOOR_DIV gt for op=trt.ElementWiseOperation.GREATER lt for op=trt.ElementWiseOperation.LESS op_and for op=trt.ElementWiseOperation.AND op_or for op=trt.ElementWiseOperation.OR eq for op=trt.ElementWiseOperation.EQUAL minimum for op=trt.ElementWiseOperation.MIN maximum for op=trt.ElementWiseOperation.MAX pow for op=trt.ElementWiseOperation.POW

它使用 TensorRT 中的 IElementWiseLayer 实现。

参数:

left – Union[Tensor, int, float] 第一个输入。如果该输入是整数或浮点数，则该函数会创建一个常量张量。
right – Union[Tensor, int, float] 第二个输入。如果该输入是整数或浮点数，则该函数会创建一个常量张量。
op – trt.ElementWiseOperation 要执行的二元运算。

返回值:

此逐元素运算生成的张量。

tensorrt_llm.functional.outer( input: Tensor, vec2: Tensor, ) → Tensor[source]#

添加一个操作来计算两个张量之间的外积。

该操作创建一个 Einsum 节点。

参数:

input – Tensor 第一个输入张量。
vec2 – Tensor 第二个输入张量。

返回值:

此层生成的输出张量。

tensorrt_llm.functional.pad( input: Tensor, pad: Sequence[int] | Tensor, mode: str = 'constant', value: float | None = None, ) → Tensor[source]#

添加一个填充层。

填充层在输入张量的开头和结尾添加零填充。从最后一个维度开始并向前移动来描述用于填充输入某些维度的大小。

输入张量的 [len(pad) / 2] 个维度将被填充。例如，要仅填充输入张量的最后一个维度，则 pad 的形式为 [padding_left, padding_right]；要填充输入张量的最后 2 个维度，则使用 [padding_left, padding_right, padding_top, padding_bottom]；要填充最后 3 个维度，则使用 [padding_left, padding_right, padding_top, padding_bottom, padding_front, padding_back]。

参数:

input – Tensor 要执行 padding_2d 的输入张量。
pad – sequence of int 用于填充的 m 个元素的元组，其长度 m 满足 m <= 2*输入维度，并且 m 是偶数。
mode – str 仅支持“constant”。
value – float “constant”填充的填充值。默认值：0。

返回值:

插入层生成的张量。

tensorrt_llm.functional.permute( input: Tensor, dims: Sequence[int], ) → Tensor[source]#

添加一个操作来置换张量的维度。

输入张量的维度根据“dims”中的维度序列进行置换。该操作映射到 tensorrt.IShuffleLayer，其中第二次转置由“dims”中的索引描述。

给定一个秩为 N 的张量，置换的结果是一个秩为 N 的张量，其中第 i 个输入维度映射到第 dims[i] 个维度。

例如，permute(input, [1, 0]) 将通过置换行和列来转置 2D 张量。

参数:

input – Tensor 要置换的输入张量。
dims – Sequence[int] 置换的描述。

返回值:

置换层生成的张量。

tensorrt_llm.functional.pow( left: ~tensorrt_llm.functional.Tensor | int | float, right: ~tensorrt_llm.functional.Tensor | int | float, *, op: ~tensorrt.tensorrt.ElementWiseOperation = <ElementWiseOperation.POW: 6>, ) → Tensor#

添加具有两个输入的逐元素运算。

以下闭包在 functional.* 中定义。

add for op=trt.ElementWiseOperation.SUM sub for op=trt.ElementWiseOperation.SUB mul for op=trt.ElementWiseOperation.PROD div for op=trt.ElementWiseOperation.DIV floordiv for op=trt.ElementWiseOperation.FLOOR_DIV gt for op=trt.ElementWiseOperation.GREATER lt for op=trt.ElementWiseOperation.LESS op_and for op=trt.ElementWiseOperation.AND op_or for op=trt.ElementWiseOperation.OR eq for op=trt.ElementWiseOperation.EQUAL minimum for op=trt.ElementWiseOperation.MIN maximum for op=trt.ElementWiseOperation.MAX pow for op=trt.ElementWiseOperation.POW

它使用 TensorRT 中的 IElementWiseLayer 实现。

参数:

left – Union[Tensor, int, float] 第一个输入。如果该输入是整数或浮点数，则该函数会创建一个常量张量。
right – Union[Tensor, int, float] 第二个输入。如果该输入是整数或浮点数，则该函数会创建一个常量张量。
op – trt.ElementWiseOperation 要执行的二元运算。

返回值:

此逐元素运算生成的张量。

tensorrt_llm.functional.prod(input: ~tensorrt_llm.functional.Tensor, *, op: ~tensorrt.tensorrt.ReduceOperation = <ReduceOperation.PROD: 1>, dim: int | ~typing.Tuple[int], keepdim: bool = False) → Tensor#

添加一个沿维度进行的归约操作。

它使用 TensorRT 中的 IReduceLayer 实现。

参数:

input – Tensor 输入张量。
op – trt.ReduceOperation 要执行的归约操作。选项：SUM、PROD、MAX、MIN、AVG
dim – int 执行归约的维度。
keepdim – bool 该维度是否保留在缩减的张量中？如果为 True，则保留该维度，否则将其从形状中删除。

返回值:

此缩减操作生成的张量。

tensorrt_llm.functional.quick_gelu( x: Tensor, ) → Tensor[source]#

tensorrt_llm.functional.rand( shape: Tensor, low: float = 0, high: float = 1, dtype: str | DataType = 'float32', ) → Tensor[source]#

此操作添加一个填充层，该层生成具有指定形状和数据类型的随机（均匀）张量。

参数:

shape – Tensor 需要生成的张量的形状。
low – float 用于随机数的范围的最小值（包括）。
high – float 用于随机数的范围的最大值（包括）。
dtype – Union[str, trt.DataType] 输出张量所需的數據類型。

返回值:

填充层生成的随机张量。

tensorrt_llm.functional.rearrange(

inputs: Tensor | Sequence[Tensor],

expression: str,

**kwargs,

) → Tensor[source]#

在张量上添加重新排列操作。

此操作是针对多维张量的一种易于阅读的智能元素重新排序，包括转置（轴置换）、重塑（视图）、挤压、取消挤压、堆叠、连接和其他操作的功能。请参阅：https://einops.rocks/api/rearrange/

例如，如果输入张量的形状为 [32, 30, 40, 3]，并运行: rearrange(x, ‘b (h h1) (w w1) c -> b h w 1 (c h1 w1) 1’, h1=2, w1=2)

它将生成一个形状为 [32, 15, 20, 1, 12, 1] 的张量。

参数:

input – Union[Tensor, Sequence[Tensor]] 如果它是张量，它将直接对其进行操作。否则，如果它是一个序列，它将把它连接成一个张量，然后对其进行操作。
expression – str 关于如何以易于阅读的方式重新排序张量的表达式。
kwargs – 用于使用特定值设置某些标识符的关键字参数。

返回值:

此操作的输出张量。

tensorrt_llm.functional.recv( tensor: Tensor, src: int, ) → Tensor[source]#

添加一个操作，该操作执行从另一个秩接收到秩的 recv。

recv 操作从另一个秩上的秩接收张量。如果秩“i”从秩“j”接收张量，则秩“j”必须具有对应的“send”操作到秩“j”。参见“send”。

该操作使用包装 NCCL 点对点 recv 操作的插件来实现。请参阅 https://docs.nvda.net.cn/deeplearning/nccl/user-guide/docs/api/p2p.html#ncclrecv 了解详情。

参数:

tensor – Tensor 输入张量。
src – int 将张量发送到的秩。

返回值:

该层生成的张量。

tensorrt_llm.functional.reduce( input: Tensor, op: ReduceOperation, dim: int | Tuple[int], keepdim: bool = False, ) → Tensor[source]#

添加一个沿维度进行的归约操作。

它使用 TensorRT 中的 IReduceLayer 实现。

参数:

input – Tensor 输入张量。
op – trt.ReduceOperation 要执行的归约操作。选项：SUM、PROD、MAX、MIN、AVG
dim – int 执行归约的维度。
keepdim – bool 该维度是否保留在缩减的张量中？如果为 True，则保留该维度，否则将其从形状中删除。

返回值:

此缩减操作生成的张量。

tensorrt_llm.functional.reduce_scatter( tensor: Tensor, group: List[int], ) → Tensor[source]#

tensorrt_llm.functional.relu( input: ~tensorrt_llm.functional.Tensor, *, act_type: ~tensorrt.tensorrt.ActivationType = <ActivationType.RELU: 0>, ) → Tensor#

添加激活函数。

参数:

input – Tensor 应用激活函数的输入张量。
act_type – trt.ActivationType 激活的类型（RELU、TANH、SIGMOID、…）。

以下闭包在 functional.* 中定义。

relu for op=trt.ActivationType.RELU tanh for op=trt.ActivationType.TANH sigmoid for op=trt.ActivationType.SIGMOID

返回值:: 激活层生成的张量。

tensorrt_llm.functional.repeat( input: Tensor, sizes: Sequence[int], ) → Tensor[source]#

沿指定维度重复张量。

参数:

input – Tensor 要重复的张量。
sizes – Sequence[int] 沿每个维度重复张量的次数。

返回值:

一个张量，除了沿指定维度重复的输入张量外。

tensorrt_llm.functional.repeat_interleave( tensor: Tensor, repeats: int, dim: int, ) → Tensor[source]#

沿轴重复张量的元素。

参数:

repeats – int 沿指定轴的重复次数。
dim – int 执行重复操作的维度。

返回值:

一个与输入具有相同形状的张量，除了沿指定维度重复的元素外。

TODO：允许 repeats 是整数列表，并且 dim 是未指定的。

tensorrt_llm.functional.rg_lru( input: Tensor, A: Tensor, state_or_ptr: Tensor, host_request_types: Tensor, last_token_ids: Tensor, dim: int, dtype: str, block_size: int = 0, y: Tensor | None = None, y_bias: Tensor | None = None, gate: Tensor | None = None, gate_bias: Tensor | None = None, gate_x: Tensor | None = None, gate_x_bias: Tensor | None = None, gate_a: Tensor | None = None, gate_a_bias: Tensor | None = None, slot_mapping: Tensor | None = None, )[source]#

参数:

input – Tensor (On GPU) 输入张量。它的形状是 [batch_size, seq_len, dim]
A – Tensor (On GPU) 一个矩阵。它的形状是 [dim]
state_or_ptr – Tensor (On GPU or CPU) lru 状态张量。它的形状是 [batch_size, dstate, dim] 或形状为 [1] 的 CPU 张量，用于分页状态的指针。
host_request_types – Tensor（在 CPU 上）主机上的张量，指示请求是在上下文还是生成阶段。其形状为 [batch_size]。请参阅 docs/source/advanced/gpt-attention.md 中的 Inflight Batching。
last_token_ids – Tensor（在 GPU 上）批次中序列长度的包含前缀和或序列的长度。
dim – int RG_LRU 块的内部维度
block_size – int 块对角线性层的块大小。它用于支持启用融合门控的情况。
dtype – str 数据类型
y – Tensor (On GPU)（可选）y 张量。它的形状是 [batch_size, seq_len, dim]
y_bias – Tensor (On GPU)（可选）y_bias 张量。它的形状是 [dim]。如果 y_bias 不是 None，我们将在该函数中融合 GELU(y + y_bias)。
gate – Tensor (On GPU)（可选）gate 张量。它的形状是 [batch_size, seq_len, 2 * dim]。如果 gate 不是 None，我们将融合 gate_x 和 gate_a，否则使用这两个张量。
gate_bias – Tensor (On GPU)（可选）gate_bias 张量。它的形状是 [2 * block_num, dim // block_num]。如果 gate_bias 不是 None，我们将在此函数中融合偏置添加。
gate_x – Tensor (On GPU)（可选）gate_x 张量。它的形状是 [batch_size, seq_len, dim]
gate_x_bias – Tensor (On GPU)（可选）gate_x_bias 张量。它的形状是 [block_num, dim // block_num]。如果 gate_x_bias 不是 None，我们将在此函数中融合偏置添加。
gate_a – Tensor (On GPU)（可选）gate_a 张量。它的形状是 [batch_size, seq_len, dim]
gate_a_bias – Tensor (On GPU)（可选）gate_a_bias 张量。它的形状是 [block_num, dim // block_num]。如果 gate_a_bias 不是 None，我们将在此函数中融合偏置添加。
slot_mapping – Tensor (On GPU)（可选）状态中的实际页面索引。它的形状是 [dim]，用于分页状态，每个页面的形状是 [dstate, dim]

tensorrt_llm.functional.rms_norm( input: Tensor, normalized_shape: int | Tuple[int], num_groups: int = 1, weight: Tensor | None = None, eps: float = 1e-06, ) → Tensor[source]#

在张量上添加 RMS 归一化操作。

该操作将其输入张量应用 rms 归一化。在其最简单的形式中，对于大型语言模型，“normalized_shape”应设置为激活张量的隐藏维度。否则，它是张量的归一化部分的形状（从最右边的维度开始）。

“weight”张量对应于 rms-norm 公式中的“gamma”。“eps”值在计算平方根之前添加到方差中。

参数:

input – Tensor 要归一化的张量。
normalized_shape – Union[int, Tuple[int]] 归一化的子张量的形状。使用“hidden_dim”来归一化 LLM 中激活张量的最内层维度。
num_groups – int = 1 组大小。
weight – Optional[Tensor] = None 层归一化中的 “gamma” 项。它的形状必须是 ‘normalized_shape’。
eps – float 要添加到平方根中的方差的 epsilon 项。weig

返回值:

该操作的输出张量。

tensorrt_llm.functional.round( input: ~tensorrt_llm.functional.Tensor, *, op: ~tensorrt.tensorrt.UnaryOperation = <UnaryOperation.ROUND: 22>, ) → Tensor#

对单个输入添加逐元素运算。

以下闭包在 functional.* 中定义。

round for op=trt.UnaryOperation.ROUND sqrt for op=trt.UnaryOperation.SQRT exp for op=trt.UnaryOperation.EXP sin for op=trt.UnaryOperation.SIN cos for op=trt.UnaryOperation.COS abs for op=trt.UnaryOperation.ABS log for op=trt.UnaryOperation.LOG

它使用 TensorRT 中的 IUnaryLayer 实现。

参数:

input – Tensor 输入张量。
op – trt.UnaryOperation 要执行的一元运算。

返回值:

此逐元素运算生成的张量。

tensorrt_llm.functional.scatter( input: Tensor, dim: int, indices: Tensor, updates: Tensor, ) → Tensor[source]#

此操作添加一个层，该层通过按元素复制输入张量的值来创建输出张量，然后通过给定的

indices 和 updates 张量更新值。对于 2D 输入张量，它首先将输入复制到输出，然后像以下一样更新输出张量 updates 中的每个条目

output[indices[i][j]][j] = updates[i][j] 如果 dim=0 output[i][indices[i][j]] = updates[i][j] 如果 dim=1

如果 input 张量为 [[1, 2, 3], [4, 5, 6]]，indices 张量为 [[1, 2], [0, 1]]，updates 张量为 [[-1, -2], [-3, -4]] 且 dim=1，则输出张量将为 [[1, -1, -2], [-3, -4, 6]]。参数

input: Tensor
需要更新的输入数据。

dim: int
要执行 scatter 的轴。

indices: Tensor
与输入具有相同秩的整数张量，指示要更新的位置。

updates: Tensor
与 indices 张量具有相同形状的数据张量，包含更新值。

返回值
由元素 scatter 层创建的张量。

tensorrt_llm.functional.scatter_nd( input: Tensor, mask: Tensor, source: Tensor, ) → Tensor[source]#

Scatter_nd 是一种张量操作，它基于索引在张量中写入或更新值。

参数:

input – Tensor 要更新的输入张量
mask – Tensor 一个索引张量，指定要更新的 data 中的位置。
source – Tensor 一个值张量，要写入或 scatter 到 data 中。

返回值:

新的张量，具有与输入张量数据相同的形状，源张量的值根据掩码张量指定的的位置分散或写入到输出张量中。

tensorrt_llm.functional.select( input: Tensor, dim: int, index: Tensor | int, ) → Tensor[source]#

添加一个操作，从张量中选择一部分元素。

给定一个输入张量，该函数创建一个操作，选择维度“dim”中第index个切片的元素，以创建一个新的张量。输出张量的形状中，输入维度“dim”被移除。

“index”可以是整数或包含单个元素的一维张量。

例如，对于输入=[[4, 2, 5], [2, 1, 2], [4, 7, 1]]，其形状为 [3, 3]，

select(input, 0, 1)

将创建一个形状为[3]的张量，其中包含[2, 1, 2]。

关于输出张量的形状，维度“dim”被移除。这意味着对于形状为[4, 2, 6, 3]的张量，

select(input, 2, 4)

将从第3个维度（dim == 2）选择第5个切片（index == 4），并返回形状为[4, 2, 3]的张量（即第3个维度被移除）。

该操作映射到 TensorRT IGatherLayer。

参数:

input – Tensor 要从中选择的输入张量。
dim – int 要从中选择的维度。
index – Union[Tensor, int] 要选择的“dim”维度中的切片索引。

返回值:

包含所选切片的张量。

tensorrt_llm.functional.selective_scan( input: Tensor, state_or_ptr: Tensor, delta: Tensor, delta_bias: Tensor, A: Tensor, BC: Tensor, D: Tensor, host_request_types: Tensor, last_token_ids: Tensor, dim: int, dstate: int, dt_rank: int, delta_softplus: bool, dtype: str, z: Tensor | None = None, host_context_lengths: Tensor | None = None, slot_mapping: Tensor | None = None, nheads: int = 1, ngroups: int = 1, chunk_size: int = 256, mamba_version: str = 'Mamba1', )[source]#

参数:

input – Tensor (On GPU) 输入张量。它的形状是 [batch_size, seq_len, dim]
state_or_ptr – 张量 (在 GPU 或 CPU 上) SSM 状态张量。它的形状是 [batch_size, dstate, dim]。或者形状为 [1] 的 CPU 张量，用于分页状态的指针。
delta – 张量（在 GPU 上）delta 张量。mamba：其形状为 [batch_size, seq_len, dim] 或 [num_tokens, dim] (用于 remove_input_padding)。mamba2：其形状为 [batch_size, seq_len, nheads] 或 [num_tokens, nheads] (用于 remove_input_padding)
delta_bias – 张量（在 GPU 上）delta 偏差张量。mamba：其形状为 [dim]。mamba2：其形状为 [nheads]
A – 张量（在 GPU 上）A 矩阵。mamba：其形状为 [dstate, dim]。mamba2：其形状为 [nheads]
BC – 张量（在 GPU 上）B 和 C 矩阵。mamba：其形状为 [batch_size, seq_len, dstate * 2] 或 [num_tokens, dstate * 2] (用于 remove_input_padding)。mamba2：其形状为 [batch_size, seq_len, ngroups * dstate * 2] 或 [num_tokens, ngroups * dstate * 2] (用于 remove_input_padding)
D – 张量（在 GPU 上）D 矩阵。mamba：其形状为 [dim]。mamba2：其形状为 [nheads]
host_request_types – 张量（在 CPU 上）主机上的张量，指示请求是否处于上下文或生成阶段。它的形状是 [batch_size]。请参阅 docs/source/advanced/gpt-attention.md 中的 Inflight Batching
last_token_ids – Tensor（在 GPU 上）批次中序列长度的包含前缀和或序列的长度。
dim – int SSM 块的内部维度
dstate – int SSM 块的状态维度
dt_rank – int dt_proj 的秩维度
delta_softplus – bool 是否将 softplus 应用于 delta。
dtype – str 数据类型
z – 张量（在 GPU 上）（可选）z 张量。其形状为 [batch_size, seq_len, dim] 或 [num_tokens, dim] (用于 remove_input_padding)
host_context_lengths – Tensor (在 CPU 上) (可选) 包含不同输入长度的主机张量。
slot_mapping – Tensor (On GPU)（可选）状态中的实际页面索引。它的形状是 [dim]，用于分页状态，每个页面的形状是 [dstate, dim]
nheads – int（可选）头的数量。
ngroups – int（可选）组的数量。
chunk_size – int（可选）chunk_size 用于 chunk_scan 内核。
mamba_version – int（可选）Mamba 版本，默认支持 Mamba1。

tensorrt_llm.functional.send( tensor: Tensor, tgt: int, ) → Tensor[source]#

添加一个操作，该操作执行从一个 rank 到另一个 rank 的发送。

发送操作将张量从一个 rank 发送到另一个 rank。如果 rank ‘i’ 将张量发送到 rank ‘j’，则 rank ‘j’ 必须具有来自 rank ‘i’ 的相应“recv”操作。见“recv”。

该操作是使用一个插件实现的，该插件包装了 NCCL 点对点发送操作。有关详细信息，请参阅 https://docs.nvda.net.cn/deeplearning/nccl/user-guide/docs/api/p2p.html#ncclsend。

参数:

tensor – Tensor 输入张量。
tgt – int 接收张量的 rank。

返回值:

该层生成的张量。

tensorrt_llm.functional.shape( input: Tensor, dim: int | None = None, cast_to_dtype: str | DataType | None = None, clip_before_cast: Sequence[int] = None, ) → Tensor[source]#

添加一个创建形状张量的操作。

形状张量可以是输入张量的形状（当参数 dim 为 None 时），也可以是对应于第 dim 个维度大小的标量（秩为 0 的张量）。

参数:

input – Tensor 我们要从中提取形状或一个维度大小的输入张量。
dim – Optional[int] 要提取大小的维度。如果为 None，则返回输入张量的整个形状。

返回值:

包含输入张量的形状（如果“dim”为 None）或输入张量的维度“dim”中的大小的张量。如果“dim”为“None”，则该张量与输入张量具有相同的秩，否则其秩为 0。

tensorrt_llm.functional.sigmoid( input: ~tensorrt_llm.functional.Tensor, *, act_type: ~tensorrt.tensorrt.ActivationType = <ActivationType.SIGMOID: 1>, ) → Tensor#

添加激活函数。

参数:

input – Tensor 应用激活函数的输入张量。
act_type – trt.ActivationType 激活的类型（RELU、TANH、SIGMOID、…）。

以下闭包在 functional.* 中定义。

relu for op=trt.ActivationType.RELU tanh for op=trt.ActivationType.TANH sigmoid for op=trt.ActivationType.SIGMOID

返回值:: 激活层生成的张量。

tensorrt_llm.functional.silu( input: Tensor, ) → Tensor[source]#

添加 SiLU (x * sigmoid(x)) 操作。

参数:: input – Tensor 应用激活函数的输入张量。
返回值:: 激活层生成的张量。

tensorrt_llm.functional.sin( input: ~tensorrt_llm.functional.Tensor, *, op: ~tensorrt.tensorrt.UnaryOperation = <UnaryOperation.SIN: 6>, ) → Tensor#

对单个输入添加逐元素运算。

以下闭包在 functional.* 中定义。

round for op=trt.UnaryOperation.ROUND sqrt for op=trt.UnaryOperation.SQRT exp for op=trt.UnaryOperation.EXP sin for op=trt.UnaryOperation.SIN cos for op=trt.UnaryOperation.COS abs for op=trt.UnaryOperation.ABS log for op=trt.UnaryOperation.LOG

它使用 TensorRT 中的 IUnaryLayer 实现。

参数:

input – Tensor 输入张量。
op – trt.UnaryOperation 要执行的一元运算。

返回值:

此逐元素运算生成的张量。

tensorrt_llm.functional.slice( input: Tensor, starts: Tensor | Sequence[int], sizes: Tensor | Sequence[int], strides: Tensor | Sequence[int] = None, mode: SampleMode = None, fill_value: float | Tensor = None, ) → Tensor[source]#

添加一个从张量中提取切片的操作。

正如 TensorRT 文档中 ISliceLayer 的描述，切片层有两个变体：静态和动态。

对于静态切片，此函数在图层创建时，通过整数序列获取不同维度上的 starts 和 sizes 值来进行切片。对于动态切片，它接受 starts 和 sizes 作为 tensorrt.ITensor。

切片层为每个维度选择输入张量中的起始位置，并使用步长 1 跨输入张量将元素复制到输出张量。如果未指定为整数序列，则 Start 和 Size 张量必须是 1-D int32 形状张量。

例如，对于 input = [[0, 2, 4], [1, 3, 5]]，调用

slice(input, start=[1, 0], size=[1, 2])

将产生张量 [[1, 3]] 作为输出。当 TensorRT 执行切片运算符时，它将复制一行（因为 size[0] == 1），从第二行开始（因为 start[0] == 1），并复制两列（size[1] == 2），从第一列开始（因为 start[1] == 0）。

在伪代码中，对于 2D 张量（可以很容易地扩展到更多维度），该操作的行为可以描述如下

output = Tensor(shape=sizes) for ii in range(sizes[0])

for jj in range(sizes[1])
output[ii][jj] = input[starts[0]+ii][starts[1]+jj]

请注意，在深度学习框架中，使用范围 [start:end] 进行类似操作是很常见的。可以通过设置 sizes 参数来模拟它，以便在每个维度中 [start:start+size] == [start:end]，即 size = end-start。

TensorRT 支持不同的切片模式，但该函数将该选择限制为 mode == tensorrt.SampleMode.STRICT_BOUNDS。

参数:

input – Tensor 要执行切片的输入张量。
starts – Union[Tensor, Sequence[int]] 输入张量中每个维度的起始点。
sizes – Union[Tensor, Sequence[int]] 切片张量（输出）的每个维度中的元素数量。
strides – Union[Tensor, Sequence[int]] 从输入张量中的起始点获取的步长。
mode – trt.SampleMode 控制切片操作如何处理越界坐标的模式。

返回值:

切片层生成的张量。

tensorrt_llm.functional.softmax( input: Tensor, dim: int | None = None, ) → Tensor[source]#

添加一个在张量上计算 softmax 的操作。

如果在维度 ‘dim’ 中指定，则该操作在输入张量的维度 ‘dim’ 上计算 softmax。否则，它应用于最后一个维度。

它将 ISoftmaxLayer 插入到 TensorRT 图中。

参数:

input – Tensor 要应用 softmax 的输入张量。
dim – Optional[int] 用于应用 softmax 的维度。

返回值:

softmax 层的输出张量。

tensorrt_llm.functional.softplus( input: Tensor, beta: float, threshold: float, ) → Tensor[source]#

根据 PyTorch 定义添加 softplus 激活。

有关该函数的说明，请参见 https://pytorch.ac.cn/docs/stable/generated/torch.nn.functional.softplus.html#torch-nn-functional-softplus。

参数:

input – Tensor 输入 TensorRT-LLM 张量。
beta – float 用于 softplus 计算的参数。
threshold – float 当 input * beta > threshold 时恢复为线性函数的阈值

返回值:

该层创建的输出张量。

tensorrt_llm.functional.split( tensor: Tensor, split_size_or_sections: int | Sequence[int], dim: int = 0, ) → Sequence[Tensor][source]#

添加将张量分割成子张量的操作。

此操作创建一个张量列表，这些张量是通过沿维度 ‘dim’ 切割输入张量而获得的。如果 ‘split_size_or_sections’ 是一个整数，则张量被拆分为 ‘input.shape[dim] / split_size_or_sections’ 个切片。如果 ‘split_size_or_sections’ 是大小列表，则张量被拆分为 ‘len(split_size_or_sections)’ 个切片，并且第 i 个切片的大小由 ‘split_size_or_sections[i]’ 给出。

当前实现存在一些限制

输入张量必须是静态的（没有动态维度）。

如果 ‘split_size_or_sections’ 是一个整数，则输入的 ‘dim’ 维度中的元素数量必须是 ‘split_size_or_sections’ 的倍数：‘input.shape[dim] % split_size_or_sections == 0’。

如果 ‘split_size_or_sections’ 是一个序列，则 ‘split_size_or_sections’ 中元素的总和必须等于 ‘dim’ 维度中的大小：‘input.shape[dim] == sum(ii for ii in split_size_or_sections)’。

该操作使用每个输出切片的 ‘slice’ 操作来实现。

参数:

tensor – Tensor 要切片的输入张量。
split_size_or_sections – Union[int, Sequence[int]] 如果它是一个整数，它会编码每个切片的大小。否则，如果它是一个序列，它是每个切片的大小。
dim – int 要切片的张量的维度。

返回值:

由不同操作产生的张量列表。

tensorrt_llm.functional.sqrt( input: ~tensorrt_llm.functional.Tensor, *, op: ~tensorrt.tensorrt.UnaryOperation = <UnaryOperation.SQRT: 2>, ) → Tensor#

对单个输入添加逐元素运算。

以下闭包在 functional.* 中定义。

round for op=trt.UnaryOperation.ROUND sqrt for op=trt.UnaryOperation.SQRT exp for op=trt.UnaryOperation.EXP sin for op=trt.UnaryOperation.SIN cos for op=trt.UnaryOperation.COS abs for op=trt.UnaryOperation.ABS log for op=trt.UnaryOperation.LOG

它使用 TensorRT 中的 IUnaryLayer 实现。

参数:

input – Tensor 输入张量。
op – trt.UnaryOperation 要执行的一元运算。

返回值:

此逐元素运算生成的张量。

tensorrt_llm.functional.squared_relu( x: Tensor, ) → Tensor[source]#

添加一个平方 ReLU 操作。

此函数应用 ReLU 并对输出求平方。

参数:: input – Tensor 应用激活函数的输入张量。
返回值:: 激活层生成的张量。

tensorrt_llm.functional.squeeze( input: Tensor, dim: int | Sequence[int] | None = None, zero_is_placeholder: bool = False, )[source]#

添加一个从张量中删除单例维度的操作。

此函数创建一个操作，该操作删除输入张量中位置 ‘dim’ 的单例维度（大小为 1 的维度）。它适用于 ‘dim’ 的负值。

例如，对于形状为 [1, 4, 1, 4] 的张量 ‘input’

squeeze(input, 0) 将产生形状为 [4, 1, 4] 的输出，squeeze(input, 2) 将产生形状为 [1, 4, 4] 的输出，squeeze(input, [0, 2]) 将产生形状为 [4, 4] 的输出，squeeze(input, [-2]) 将产生形状为 [1, 4, 4] 的输出，

参数:

input – Tensor 要从中删除单例维度的输入张量。
dim – Union[int, Sequence[int]] 输入张量中单例维度的索引。

返回值:

该层产生的张量。

tensorrt_llm.functional.stack( inputs: Sequence[Tensor], dim: int = 0, ) → Tensor[源代码]#

添加一个沿新维度连接输入张量的操作。

该函数创建一个操作，为所有输入张量创建一个新维度，然后沿着该新维度连接它们。

“inputs”中的所有张量必须具有相同的形状。

for ii in range(inputs[0].rank())
assert all(inp.shape[ii] == inputs[0].shape[ii] for inp in inputs)

输出张量的形状定义为

output.rank() = inputs[0].rank() + 1

output.shape[dim] = len(inputs)

for ii in range(inputs[0].rank())

if ii < dim
output.shape[ii] = inputs[0].shape[ii]

else
output.shape[ii+1] = inputs[0].shape[ii]

例如，给定一个由两个 2D 张量 [[0, 1], [2, 3]] 和 [[4, 5], [6, 7]] 组成的序列，它们的形状均为 [2, 2]，

stack(inputs, 0)

将生成 [[[0, 1], [2, 3]], [[4, 5], [6, 7]]]，形状为 [2, 2, 2]，并且

stack(inputs, 1)

将生成 [[[0, 1], [4, 5]], [[2, 3], [6, 7]]]，形状为 [2, 2, 2]。

参数

inputsSequence[Tensor]
要堆叠的张量序列。

dimint
执行堆叠的维度。

返回值
一个张量，包含沿新维度堆叠的输入张量。

tensorrt_llm.functional.sub( left: ~tensorrt_llm.functional.Tensor | int | float, right: ~tensorrt_llm.functional.Tensor | int | float, *, op: ~tensorrt.tensorrt.ElementWiseOperation = <ElementWiseOperation.SUB: 4>, ) → Tensor#

添加具有两个输入的逐元素运算。

以下闭包在 functional.* 中定义。

add for op=trt.ElementWiseOperation.SUM sub for op=trt.ElementWiseOperation.SUB mul for op=trt.ElementWiseOperation.PROD div for op=trt.ElementWiseOperation.DIV floordiv for op=trt.ElementWiseOperation.FLOOR_DIV gt for op=trt.ElementWiseOperation.GREATER lt for op=trt.ElementWiseOperation.LESS op_and for op=trt.ElementWiseOperation.AND op_or for op=trt.ElementWiseOperation.OR eq for op=trt.ElementWiseOperation.EQUAL minimum for op=trt.ElementWiseOperation.MIN maximum for op=trt.ElementWiseOperation.MAX pow for op=trt.ElementWiseOperation.POW

它使用 TensorRT 中的 IElementWiseLayer 实现。

参数:

left – Union[Tensor, int, float] 第一个输入。如果该输入是整数或浮点数，则该函数会创建一个常量张量。
right – Union[Tensor, int, float] 第二个输入。如果该输入是整数或浮点数，则该函数会创建一个常量张量。
op – trt.ElementWiseOperation 要执行的二元运算。

返回值:

此逐元素运算生成的张量。

tensorrt_llm.functional.sum( input: Tensor, dim: int, keepdim: bool = False, ) → Tensor[源代码]#

添加一个沿着维度计算总和的操作。

计算输入张量沿着维度“dim”的总和。

它使用 TensorRT 中的 IReduceLayer 实现。

参数:

input – Tensor 输入张量。
dim – int 计算平均值的维度。
keepdim – bool 该维度是否保留在缩减的张量中？如果为 True，则保留该维度，否则将其从形状中删除。

返回值:

此缩减操作生成的张量。

tensorrt_llm.functional.swiglu( input: Tensor, ) → Tensor[源代码]#

添加一个 SwiGLU (x * SiLU(gate)) 操作。

该函数接受一个张量，沿着最后一个维度将其分成两半，将 SiLU 应用于后半部分，然后将结果相乘。如果最后一个维度不是偶数，则行为未定义。

参数:: input – Tensor 应用激活函数的输入张量。
返回值:: 激活层生成的张量。

tensorrt_llm.functional.tanh( input: ~tensorrt_llm.functional.Tensor, *, act_type: ~tensorrt.tensorrt.ActivationType = <ActivationType.TANH: 2>, ) → Tensor#

添加激活函数。

参数:

input – Tensor 应用激活函数的输入张量。
act_type – trt.ActivationType 激活的类型（RELU、TANH、SIGMOID、…）。

以下闭包在 functional.* 中定义。

relu for op=trt.ActivationType.RELU tanh for op=trt.ActivationType.TANH sigmoid for op=trt.ActivationType.SIGMOID

返回值:: 激活层生成的张量。

tensorrt_llm.functional.topk( input: Tensor, k: Tensor | int, dim: int, largest: bool = True, prefer_plugin: bool = True, ) → Tuple[Tensor, Tensor][源代码]#

添加一个 topk 操作。

正如 ONNX 文档中所解释的，

onnx/onnx

注意：与 ONNX topk 操作的一个区别是，TensorRT 层始终对输出进行排序。

检索指定轴上最大的前 K 个元素。给定一个形状为 [a_1, a_2, …, a_n, r] 的输入张量和整数参数 k，返回两个输出：形状为 [a_1, a_2, …, a_{axis-1}, k, a_{axis+1}, … a_n] 的值张量，其中包含指定轴上最大的前 k 个元素的值。形状为 [a_1, a_2, …, a_{axis-1}, k, a_{axis+1}, … a_n] 的索引张量，其中包含最大的前 k 个元素的索引（来自输入张量的原始索引）。

参数:

input – Tensor 输入张量。
k – int 对应于要检索的顶部元素数量的单个正值
dim – int 用于计算 topk 索引的维度。
largest – bool 控制是返回最大还是最小元素
prefer_plugin – bool 如果 dim 是最后一个维度且 k 是静态的，是否使用 topkLastDim 插件。

返回值:

此 topk 操作生成的张量（值，索引）。

tensorrt_llm.functional.transpose( input: Tensor, dim0: int, dim1: int, ) → Tensor[源代码]#

添加一个操作来转置张量的两个维度。

该操作生成一个张量，其中维度“dim0”和“dim1”被置换。如果张量的秩大于 2，则其他维度保持不变。

该函数是在“functional.permute”函数上构建的辅助函数。

参数:

input – Tensor 要转置的输入张量。
dim0 – int 要转置的第一个维度。
dim1 – int 要转置的第二个维度。

返回值:

置换层生成的张量。

tensorrt_llm.functional.unary( input: Tensor, op: UnaryOperation, ) → Tensor[源代码]#

对单个输入添加逐元素运算。

以下闭包在 functional.* 中定义。

round for op=trt.UnaryOperation.ROUND sqrt for op=trt.UnaryOperation.SQRT exp for op=trt.UnaryOperation.EXP sin for op=trt.UnaryOperation.SIN cos for op=trt.UnaryOperation.COS abs for op=trt.UnaryOperation.ABS log for op=trt.UnaryOperation.LOG

它使用 TensorRT 中的 IUnaryLayer 实现。

参数:

input – Tensor 输入张量。
op – trt.UnaryOperation 要执行的一元运算。

返回值:

此逐元素运算生成的张量。

tensorrt_llm.functional.unbind(input: Tensor, dim: int = 0)[源代码]#

删除张量维度。

返回沿给定维度的所有切片的元组，且不包含该维度。

tensorrt_llm.functional.unsqueeze(input: Tensor, axis: int)[源代码]#

添加一个操作以将单例维度插入张量。

该函数创建一个操作，该操作在输出张量中的位置“axis”处插入一个单例维度（大小为 1 的维度）。它适用于“axis”的负值。

例如，对于形状为 [4, 4] 的张量“input”

unsqueeze(input, 0) 将生成一个形状为 [1, 4, 4] 的输出，unsqueeze(input, 1) 将生成一个形状为 [4, 1, 4] 的输出，unsqueeze(input, -1) 将生成一个形状为 [4, 4, 1] 的输出，unsqueeze(input, -2) 将生成一个形状为 [4, 1, 4] 的输出，

参数:

input – Tensor 要使用单例维度扩展的输入张量。
axis – int 输出张量中单例维度的索引。

返回值:

该层产生的张量。

tensorrt_llm.functional.view( input: Tensor, shape: Tensor | Sequence[int], zero_is_placeholder: bool = True, ) → Tensor[源代码]#

添加一个操作以创建张量的视图。

该操作将 tensorrt.IShuffleLayer 添加到网络。如果“shape”参数是 Tensor，则该视图是动态的。否则，它是一个静态视图。

请注意，TensorRT 限制了推断维度的数量为 1。这意味着形状序列或张量不能包含超过一个 -1。此函数强制执行此约束，如果不符合，则会断言。

参数:

input – Tensor 要转置的输入张量。
shape – Union[Tensor, Sequence[int]] 新张量的形状。
zero_is_placeholder – bool 当此参数为 True 时，'shape' 中的 0 将被 'input' 中相应维度的大小替换。否则，对应于 0 的维度将被缩小。

返回值:

view/shuffle 层生成的张量。

添加一个 where (又名 select 或 if-then-else) 操作。

假设三个输入参数具有相同的形状，该函数创建操作以计算具有相同形状的张量，使得

for ii in range(mul(condition.shape))
output[ii] = left[ii] if condition[ii] else right[ii]

对于每个输入，如果条件是布尔值或者 left/right 输入是整数或浮点数，则该函数首先创建一个常量张量。然后，如果需要，它会扩展较小的张量，以确保其秩与较大的张量相同。然后，它执行选择。

它使用 TensorRT 中的 ISelectLayer 实现。

参数:

condition – Union[Tensor, bool] 条件。如果此输入是布尔值，则该函数创建一个常量张量。
left – Union[Tensor, int, float] 第一个输入。如果该输入是整数或浮点数，则该函数会创建一个常量张量。
right – Union[Tensor, int, float] 第二个输入。如果该输入是整数或浮点数，则该函数会创建一个常量张量。

返回值:

此 where 操作产生的张量。