多头注意力、多查询注意力和组查询注意力#

本文档详细介绍了 TensorRT-LLM 中类 GPT 自回归模型的多头注意力 (MHA)、多查询注意力 (MQA) 和组查询注意力 (GQA) 的实现。简单回顾一下，多头注意力是批量矩阵乘法、softmax 和另一个批量矩阵乘法的序列，如Attention Is All You Need 文章中所述。多查询注意力 (MQA) 和组查询注意力 (GQA) 是 MHA 的变体，它使用的 K/V 头数量少于查询头数量。 TensorRT-LLM、MHA、MQA 和 GQA 由运算符 tensorrt_llm.functional.gpt_attention实现。

重要提示#

如下所述，当前实现支持两种输入模式：填充和打包（非填充）。由于打包模式始终比填充模式更节省内存且速度更快，因此将来可能会删除对填充模式的支持。

填充张量和打包张量#

在 TensorRT-LLM 中，GPT 注意力运算符支持两种不同类型的 QKV 输入：填充和打包（即非填充）输入。该模式由全局配置参数 remove_input_padding 确定，该参数在 tensorrt_llm.plugin 中定义。

启用填充时（即 remove_input_padding 为 False），短于 max_sequence_length 的序列将填充到最大长度。这可能会导致过多的内存消耗以及对填充令牌的不必要计算（在围绕 MHA 块的各种矩阵乘法中）。

为了克服这个问题，TensorRT-LLM 支持一种无填充模式，其中不同的令牌被打包在一起，并且用户为运算符提供一个包含不同序列长度的 1D 张量。建议用户始终使用打包模式（并且将来可能会删除对填充模式的支持）。

上下文和生成阶段#

GPT 注意力运算符封装了 GPT 等自回归模型中上下文和生成阶段的不同实现。

上下文阶段#

如果 context_fmha_type 设置为 disabled（请参阅 tensorrt_llm.plugin），则该实现映射到一系列 GPU 内核，这些内核将在调用 softmax 运算符之前将中间 Q*K^T 张量存储在内存中。这是最慢的方法，并且内存占用量很大（与序列长度呈二次方关系）。

否则，如果 context_fmha_type 设置为 enabled 或 enabled_with_fp32_acc（第一个批量矩阵乘法中的累积被强制为 FP32），则该函数将触发一个内核，该内核使用单个内核执行 MHA/MQA 块。对于短序列，该内核使用 MHA/MQA 的基本实现。对于较长的序列，此内核使用 Flash Attention 算法，如 FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness 和 FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning 中所述。

目前，该实现会触发额外的内核来对元素应用预处理（如 RoPE）并填充 KV 缓存（请参见下文）。在未来的版本中，计划减少此类内核的数量，以提高整体性能。

FP8 上下文 FMHA#

激活 FP8 量化后，可以通过启用 FP8 上下文 FMHA ( use_fp8_context_fmha = enable) 来进一步加速注意力。

FP8 分页上下文 FMHA 也支持 fp8 量化工作流程。您需要同时指定 use_fp8_context_fmha = enable 和 use_paged_context_fmha = enable。

请注意，此功能仅在 Ada 和 Hopper 上受支持。

生成阶段#

生成阶段是使用 TensorRT-LLM 中称为masked multi-head attention的单个内核实现的。该内核能够动态地对 Q、K 和 V 元素应用预处理：添加 QKV 偏差，应用 RoPE，以及执行反量化和量化。 TensorRT-LLM 将在未来的版本中继续添加（或启用）其他功能。例如，启用对 IA3 的支持。

掩码 MHA 内核有一个特殊版本，可以在 GPU 上跨多个 CUDA 线程块分配工作，以应对 GPU 占用率低的情况。从 TRT-LLM 0.13 开始，默认启用该模式（称为多块），并且可以在运行时使用 --multi_block_mode=False 禁用它。建议用户在模型中的批量大小和头数都相对较小的情况下测试该模式。在这种情况下，小的确切定义将取决于 GPU 的模型，并且很难预测，但为了提供一个经验法则，值得在 batch_size * num_heads 小于 GPU 上的多处理器数量时测试该模式（随着进行更多研究并且软件改进，该建议可能会在将来发展）.

请注意，即使启用了多块模式，注意力运算符也不会立即触发 GPU 内核的多块版本。多块版本需要最少数量的令牌（输入 + 生成）才能比每个头使用单个 CUDA 线程块的“基本”实现更有效。它由内部启发式控制。

另一个注意事项是，由于掩码 MHA 内核使用的共享内存大小与序列长度成正比，因此在未启用多块模式时，可能存在 GPU 的共享内存不足的情况。为了使掩码 MHA 内核在这些情况下工作，强制启用多块模式并打印警告日志。

XQA 优化#

XQA 优化是 MQA/GQA 在生成阶段的另一种优化方法。它仍然是一项实验性功能，支持的配置有限。LLAMA2 70B 是它支持的模型之一。

XQA 优化支持矩阵

FP16 / BF16 计算数据类型。
FP16 / BF16 / FP8 / INT8 KV 缓存数据类型。
分页 KV 缓存 (每个块 8 / 16 / 32 / 64 / 128 个 token)。

默认情况下此功能已启用。要禁用此功能，在构建引擎时需要使用标志 --disable_xqa。请注意，还使用启发式算法来决定使用 XQA 内核还是 masked MHA 内核以获得更好的性能。这意味着即使没有设置 --disable_xqa，也可能不会使用 XQA 内核。如果您希望在可能的情况下始终使用该内核，可以设置 TRTLLM_FORCE_XQA=1 以强制使用 XQA 内核（当模型配置支持时）。有关详细的支持配置，请参阅 cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQARunner.h 中的类 DecoderXQARunner 的 shouldUse 函数。

飞行中批处理#

TensorRT-LLM 支持请求的飞行中批处理（也称为连续批处理或迭代级别批处理），以提高服务吞吐量。借助此功能，上下文阶段的序列可以与生成阶段的序列一起处理。该技术的目的是更好地交错请求以减少延迟，并更好地利用 GPU。出于效率原因 (1)，对 inflight 批处理的支持需要将输入张量打包（无填充）。

在当前的实现中，正在经历上下文阶段的序列必须在输入张量中位于生成阶段的序列之前。例如，对于序列 S0、S1 和 S2，如果 S0 和 S2 处于上下文阶段（S1 处于生成阶段），则来自 S0 和 S2 的 token 必须出现在输入张量中 S1 的 token 之前。该约束可能在未来的版本中放松或不放松。

(1) 在生成阶段填充包含单个 token 的序列至最大输入序列的长度是对资源的一种低效利用.

分块上下文#

在原始状态下，常见的行为是一次性处理所有上下文 token。此功能将上下文分成多个块。这样，上下文块可以在生成阶段与更多 token 进行批处理，这有望提高总吞吐量。分块上下文还消除了对输入长度的限制。要启用此功能，还需要启用 FMHA 分页 kv 缓存。除了最后一个块外，上下文块的大小需要是 kv 缓存块大小的整数倍。有关用法，请参阅性能最佳实践。

KV 缓存#

在生成阶段，一种常见的优化是向 MHA 内核提供一个缓存，其中包含已计算的过去 K 和 V 元素的值。该缓存称为 KV 缓存。 TensorRT-LLM 使用该技术来加速其生成阶段。在 TensorRT-LLM 中，每个 Transformer 层都有一个 KV 缓存，这意味着 KV 缓存的数量与模型中的层数一样多。当前版本的 TensorRT-LLM 支持两种不同类型的 KV 缓存： 连续的 和 分页的 KV 缓存。

连续 KV 缓存#

连续 KV 缓存是一个单体张量。它的形状是

[max_batch_size * max_beam_width, 2, num_heads, max_seqlen, hidden_dim_per_head].

当序列短于最大序列长度时，该实现会使用比所需更多的内存（即使在生成许多输出 token 后最终接近该限制，也可能需要很多步骤才能达到该点）。

分页 KV 缓存#

分页 KV 缓存将 KV 缓存分解为多个块，这些块在处理过程中由缓存管理器分发给不同的请求。该缓存管理器跟踪序列，从池中分配新块并在需要时回收这些块。有关简化的实现，请参见tensorrt_llm.runtime.KVCacheManager。更高效的 C++ 实现包含在Batch Manager中。

INT8/FP8 KV 缓存#

在其当前实现中，即使网络的其余部分以 INT8 或 FP8 运行，GPT 注意力算子也使用 FP32、FP16 和 BFloat16 输入和输出。但是，TensorRT-LLM 支持 INT8 和 FP8 (kv_cache_quant_mode=QuantMode.INT8_KV_CACHE 和 kv_cache_quant_mode=QuantMode.FP8_KV_CACHE) KV 缓存。

GPT 注意力算子会填充 KV 缓存。当启用 INT8 或 FP8 KV 缓存时，输入值必须使用缩放因子量化为 8 位。对于量化，缩放因子存储在 kv_cache_scaling_factor 张量中。它的形状是 [1]，并且当前版本仅支持每张量量化。量化使用反向比例，因为它在插件中乘以 fp_value * (1.0 / kv_cache_scaling_factor)。

在生成期间，从缓存读取的值在 MHA/MQA 内核中动态反量化，反量化可以描述为 quantized_value * kv_cache_scaling_factor。

滑动窗口注意力，循环（滚动缓冲区）KV 缓存#

TensorRT-LLM 具有一个名为 Cyclic KV Cache 的功能，该功能将 kv 缓存视为循环缓冲区。这意味着它仅存储最后 N 个 token 的 kv 缓存，其中 N 由 GenerationSession.setup 中的 max_attention_window_size 参数确定。您可以在 run.py 或 summarize.py 文件中看到此示例。当缓存已满时，新 token 的 kv 缓存将覆盖“最近最少使用”的缓存。

在上下文阶段，如果输入长度超过 max_attention_window_size，则会激活 Sliding Window Attention。这与 sliding window_size 的功能相同。

此功能有助于减少处理非常长的序列时 kv 缓存的内存占用。

该功能还支持为每层设置不同的 max_attention_window_size 值。要使用此功能，只需在使用 python 运行时会话时向 GenerationSession.setup 提供 int32 torch.Tensor 或 list，或者在使用 cpp 运行时时向 KvCacheConfig 提供一个向量。如果提供的元素数量少于层数，则提供的 tensor/list/vector 将被多次重复到层数，然后保存为新 tensor。此张量将用作 max_attention_window_size 的缓冲区，为每层设置唯一值。但是，重要的是要注意，kv 缓存的内存分配仍然依赖于缓冲区的最大值。

_请注意，循环 kv 缓存功能当前不适用于 beam searching，因为上下文 kv 缓存在 beams 之间共享。

StreamingLLM#

StreamingLLM 功能使用窗口注意力在长文本上执行高效且稳定的 LLM，这意味着只需要将 N 个 token 存储在 KV 缓存中。与 TensorRT-LLM 中的循环 KV 缓存功能类似，max_attention_window_size 参数用于确定 N。与循环 KV 缓存功能不同，前 S 个 token（称为 sink token）始终保留在注意力窗口中，其中 S 由 GenerationSession.setup 中的 sink_token_length 参数确定。但在上下文阶段，自注意力在 StreamingLLM 的官方实现中是密集的，它使用所有 token 进行计算，并且仅将 N 个 token 保存到 KV 缓存中。

此外，StreamingLLM 中还更改了相对位置嵌入。在确定相对距离并将位置信息添加到 token 时，StreamingLLM 使用缓存中的位置，而不是原始文本中的位置。

streamingllm 标志用于启用此功能。

Beam-Search#

GPT 注意力算子支持 beam-search。在上下文阶段，每个输入序列计算一个 beam。在生成阶段，MHA/MQA/GQA 内核使用一个额外的张量来重建每个 beam 的正确路径。该张量称为 cache_indirection。它的形状是 [batch_size, beam_width, max_seqlen]。

对于一个序列 si，一个束 bi 和一个 token ti，元素 cache_indirection[si][bi][ti] 是一个介于 0 和 beam_width-1 之间的整数，指示从 KV 缓存中的哪个路径读取 K 和 V 元素。此张量在采样阶段填充。

输入 QKV 张量#

输入 QKV 张量在隐藏状态的投影之后，将 Q、K 和 V 张量（沿最后一个维度连接）打包在一起。这是一个 3D 张量。 RoPE 和量化到 INT8 或 FP8（如果需要）由 GPT 注意力算子执行。

在填充模式下，它的形状是 [batch_beam_size, max_seqlen, 3 * hidden_dim]，其中 batch_beam_size 是上下文阶段的批次大小（序列的数量），以及生成阶段的批次大小乘以束宽度。不支持在填充模式下每个序列具有不同的束宽度。

在打包模式下，它的形状是 [num_tokens, 3 * hidden_dim]，其中 num_tokens 是批次中 token 的总数。对于上下文阶段的序列，序列的 token 数量对应于其输入长度（即使束宽度大于 1 以进行集束搜索）。对于生成阶段的序列，每个序列有 beam_width 个 token。每个序列的束宽度可以不同。

换句话说，计算 token 数量的伪代码是

num_tokens = 0

# Add the length of each sequence in context phase.
for seq in context_phase:
    num_tokens += seq.length

# Add the width of the beam for each sequence in generation phase.
for seq in generation_phase:
    num_tokens += seq.beam_width

旋转位置嵌入 (RoPE)#

GPT 注意力操作可以执行旋转位置嵌入 (RoPE) 的计算。当启用该操作时，rotary_embedding_dim 设置为大于 0 的值，它与其他操作融合在一起。 GPT 算子通过将 position_embedding_type 设置为 PositionEmbeddingType.rope_gpt_neox 或 PositionEmbeddingType.rope_gptj 来支持 GPT-NeoX 和 GPT-J 形式的 RoPE。

ALiBi#

GPT 注意力算子可以将 ALiBi 应用于 Q*K^T 乘积的结果。偏差是从优化内核中的 ALiBi 斜率动态计算的。

缩放因子#

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

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

交叉注意力#

除了 GPT 风格的仅解码器模型所需的作为自注意力的 MHA 之外，gpt_attention 还支持交叉注意力。

这使得 gpt_attention 可以更广泛地用作通用解码器组件。例如，Encoder-Decoder 模型使用 gpt_attention 来发布其解码器中的自注意力和交叉注意力模块。

相对注意力偏差 (RAB)#

相对注意力偏差 (RAB) 是一种相对位置建模，根据相对位置添加注意力偏差 (Q*K^T+bias)。 RAB 是一种包含相对位置信息的轻量级方法，并且用于流行的 Encoder-Decoder 模型 T5 以及 T5 系列中的其他模型。

RAB 在两种模式下受支持：i) 常规模式，用户传入提前计算的相对注意力偏差到 MHA。 ii) 隐式模式，在 MHA 中动态计算相对注意力偏差。当相对注意力偏差太大而无法放入内存时，隐式模式适用，可以通过传入 max_distance 来打开。