Skip Softmax技术在TensorRT-LLM中加速长上下文推理

通过Skip Softmax在NVIDIA TensorRT-LLM中加速长上下文推理

对于大规模部署LLM的机器学习工程师而言，随着上下文长度的增加，注意力计算的成本呈爆炸式增长，这是一个常见且严峻的挑战。无论是在检索增强生成（RAG）管道、智能体AI工作流还是生成长篇内容，注意力机制的复杂度始终是主要瓶颈。

本文介绍了一种称为Skip Softmax的技术，这是一种硬件友好、即插即用的稀疏注意力方法，可以在无需任何重新训练的情况下加速推理。接下来将了解Skip Softmax如何实现高达1.4倍的首令牌时间（TTFT）加速和高达1.4倍的每输出令牌时间（TPOT）加速，以及如何在NVIDIA TensorRT-LLM中开始使用这项技术。

Skip Softmax如何工作？

Skip Softmax的核心是提供一种动态剪枝注意力块的方法。这之所以可能，是因为它利用了Softmax函数的一个基本特性：$e^{\text{小的负数}} \approx 0$。

在标准的FlashAttention中，GPU会为查询（$Q$）和键（$K$）块计算注意力分数（logits）。然后应用softmax将这些分数归一化为概率（$P$），再与值（$V$）相乘。

然而，注意力本质上是稀疏的。对于许多注意力块，其注意力分数相对于主导令牌而言非常低，以至于它们对最终输出的贡献在统计上可以忽略不计。Skip Softmax修改了FlashAttention的内部循环，以便尽早检测到这些块并直接跳过它们。

Skip Softmax算法

该算法直接在FlashAttention内核中实现，其逻辑遵循以下启发式规则：

1. 计算局部最大值：计算当前块的最大logit值（$l_{\text{local_max}}$）。
2. 与运行最大值比较：检查当前块的局部最大值（$l{\text{local_max}}$）与运行全局最大值（$g{\text{max}}$）之间的差值是否超过一个校准后的阈值（$\Delta$）。
3. 跳过：如果满足条件，内核将跳过该块的softmax和BMM2计算，并且关键的是，跳过从高带宽内存（HBM）加载对应的$V$块。

使用Skip Softmax有哪些好处？

Skip Softmax具有即插即用的兼容性、硬件效率、灵活性和通用性。

与需要特定架构修改的方法（如线性注意力）不同，Skip Softmax与使用标准注意力机制（如MHA、GQA或MLA）的现有预训练模型兼容。它经过优化，能够充分利用NVIDIA Hopper和NVIDIA Blackwell GPU的特定张量核心和内存层次结构。它还可以与其他优化方法结合使用。例如，在预填充阶段使用XAttention，在解码阶段使用Skip Softmax，已被证明可以在不影响准确性的情况下带来显著的加速效果。

Skip Softmax的通用性在于它能同时解决预填充和解码两个阶段的瓶颈。基于在Hopper和Blackwell架构上的性能数据，Skip Softmax在带宽受限的解码阶段和计算受限的预填充阶段（尤其是在长上下文场景下）都很有益处。

带宽受限的解码阶段

在生成（解码）阶段，LLM推理通常受内存带宽限制。GPU花费更多时间移动KV缓存数据而非进行计算。

• 好处：通过尽早识别不重要的块，Skip Softmax完全避免加载相关的$V$块。
• 数据：在Llama 3.3 70B（NVIDIA GB200 NVL72）上，Skip Softmax在解码阶段预计能实现1.36倍端到端加速。

计算受限的预填充阶段

在预填充阶段（处理输入提示词时），系统受计算能力限制。

• 好处：跳过softmax和第二个矩阵乘法（BMM2）可节省大量浮点运算（FLOPs）。
• 数据：对于相同的Llama 3.3 70B模型（NVIDIA GB200 NVL72），在128K上下文长度下，预填充阶段预计可实现1.4倍的端到端加速。

长上下文场景

Skip Softmax的效果随着序列长度的增加而提高。跳过的阈值在数学上与上下文长度（$N$）相关，关系式为 $\Delta \approx -\ln N$。这意味着，随着上下文增长，安全识别和跳过高稀疏块的机会也会增加。

准确性与稀疏性之间的权衡

对于任何近似技术，一个显而易见的问题是：“这种方法如何影响准确性？”

在RULER（合成长上下文）和LongBench（现实长上下文）基准测试上的广泛测试表明，稀疏性存在一个清晰的“安全区”。

• 安全区：稀疏率在50%左右（跳过半数的块）被认为是安全区。在Llama 3.1 8B和Qwen3-8B的测试中，以约50%的稀疏率运行，在大多数任务上实现了接近无损的准确性。
• 危险区：将稀疏率推至60%以上通常会导致准确性急剧下降，尤其是在复杂的“大海捞针”式多键任务中。
• 长文本生成：对于需要生成长输出的任务（如MATH-500），Skip Softmax保持了与密集注意力相当的准确性，这与某些静态KV缓存压缩方法不同。

表 1. 与无稀疏基准相比的准确度变化

表 2. 在128K超长序列下Qwen3-30B-Instruct模型的加速效果

部署时的额外优化包括：

• 自动化校准程序，用于确定目标稀疏级别的最佳阈值。
• 稀疏感知训练，使模型对稀疏注意力模式更加鲁棒。

在NVIDIA TensorRT-LLM中开始使用Skip Softmax

Skip Softmax Attention已直接集成到NVIDIA TensorRT-LLM中，并支持在NVIDIA Hopper和NVIDIA Blackwell数据中心GPU上使用。这使得用户可以在TensorRT-LLM提供的最先进的LLM推理性能基础上，进一步加速注意力计算。

可以通过LLM API的稀疏注意力配置启用Skip Softmax Attention：

from tensorrt_llm import LLM
from tensorrt_llm.llmapi import SkipSoftmaxAttentionConfig

sparse_attention_config = SkipSoftmaxAttentionConfig(threshold_scale_factor=1000.0)
# 此外，可以为预填充和解码阶段分别配置threshold_scale_factor。
sparse_attention_config = SkipSoftmaxAttentionConfig(threshold_scale_factor={"prefill": 1000.0, "decode": 500.0})

llm = LLM(
   model="Qwen/Qwen3-30B-A3B-Instruct-2507",
   sparse_attention_config=sparse_attention_config,
   # 其他LLM参数...
)

实际阈值等于 threshold_scale_factor 除以上下文长度。

也可以通过额外的LLM API选项YAML文件指定配置。下面是一个启动OpenAI兼容端点的示例：

# 配置示例
cat >extra_llm_api_options.yaml <<EOF
sparse_attention_config:
    algorithm: skip_softmax
    threshold_scale_factor: 1000.0
EOF

# 此外，可以为预填充和解码阶段分别配置threshold_scale_factor。
cat >extra_llm_api_options.yaml <<EOF
sparse_attention_config:
    algorithm: skip_softmax
    threshold_scale_factor:
        prefill: 1000.0
        decode: 500.0
EOF

trtllm-serve Qwen/Qwen3-30B-A3B-Instruct-2507 --extra_llm_api_options extra_llm_api_options.yaml

了解更多

要了解更多信息，请参阅论文《BLASST: Dynamic Blocked Attention Sparsity via Softmax Thresholding》，以及TensorRT-LLM关于LLM API和CLI的文档。校准功能将由NVIDIA Model Optimizer支持，使用户可以指定目标稀疏度并获得所需的阈值缩放因子。

Skip Softmax稀疏注意力内核也将通过FlashInfer Python API提供。请关注即将发布的TensorRT-LLM、Model Optimizer和FlashInfer更新，以获取官方发布信息。FINISHED