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

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

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

Skip Softmax 如何工作？

Skip Softmax的核心是提供一种动态修剪注意力块的方法。这之所以可行，是因为它利用了Softmax函数的一个基本特性：exp(很小的负数) ≈ 0。

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

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

Skip Softmax 算法

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

1. 计算局部最大值：计算当前块的最大logit值（L_local_max）。
2. 与运行最大值比较：检查当前块的局部最大值（L_local_max）与运行的全局最大值（L_global_max）之间的差值是否超过一个已校准的阈值（T）。
3. 跳过：如果条件满足，内核将跳过该块的softmax和第二次矩阵乘法（BMM2）计算，并且至关重要的是，跳过从高带宽内存（HBM）加载相应的V块。

使用 Skip Softmax 有哪些好处？

Skip Softmax 提供了即插即用的兼容性、硬件效率、灵活性和通用性。

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

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

带宽受限的解码

在生成（解码）阶段，LLM推理通常受限于内存带宽。GPU移动KV缓存数据的时间比计算时间更长。

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

计算受限的预填充

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

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

长上下文场景

Skip Softmax的效果随着序列长度的增加而提高。跳过阈值（T）与上下文长度（L）在数学上存在关系：T ≈ K / L。这意味着随着上下文增长，安全识别和跳过稀疏块的机会也随之增加。

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

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

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

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

表 1. 与无稀疏性基线的准确性差异

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

部署时的其他优化包括：

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

在 TensorRT-LLM 中开始使用 Skip Softmax

Skip Softmax Attention 已直接集成到 TensorRT-LLM 中，并在某中心Hopper和某中心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的部分。校准功能将由某中心模型优化器支持，该优化器允许用户指定目标稀疏度并获得所需的阈值缩放因子。

Skip Softmax 稀疏注意力内核也将通过 FlashInfer Python API 提供。敬请关注即将发布的 TensorRT-LLM、模型优化器和 FlashInfer 更新中的官方版本。