NVIDIA 推出 Skip Softmax 提高大型语言模型推理效率

NVIDIA 推出了名为 Skip Softmax 的新技术，集成到其 TensorRT-LLM 中，承诺加速长上下文推理。NVIDIA 表示，此举是为了应对大规模部署大型语言模型（LLM）日益苛刻的计算需求而做出的回应。

理解 Skip Softmax

Skip Softmax 是一种硬件友好的插入式稀疏注意力方法，旨在提高推理速度而无需重新训练模型。它实现了最多 1.4 倍的初始标记时间（TTFT）和输出标记时间（TPOT）的加速，对于从事长篇内容生成和其他复杂 AI 工作流程的机器学习工程师来说，是一项重要创新。

Skip Softmax 的核心原理是通过利用 Softmax 函数的数学特性动态修剪注意力块。这允许尽早检测并跳过对最终输出贡献微小的注意力块，从而减少计算开销。

优势与实施

Skip Softmax 旨在与使用标准注意力机制的现有预训练模型兼容。它针对 NVIDIA 的 Hopper 和 Blackwell GPU 架构进行了优化，提供无缝的集成，提升速度和效率。值得注意的是，它可以与其他优化方法结合，例如在预填充阶段使用 XAttention 和在解码阶段使用 Skip Softmax，以实现显著的速度提升。

性能测试表明，Skip Softmax 能大幅减少解码和预填充阶段的内存带宽和计算需求。例如，在 Llama 3.3 70B 模型上，解码阶段观察到预计 1.36 倍的加速，预填充阶段在 128K 上下文长度时速度提高了 1.4 倍。

准确性与稀疏性权衡

虽然 Skip Softmax 提供了效率提高，但它也在稀疏性的“安全区”内保持准确性。各种基准测试表明，稀疏率达到 50% 时能保持几乎无损的准确性，而超过 60% 时可能导致准确性下降。这使其适用于需要生成长输出的任务，与密集注意力方法保持一致。

如何开始使用 Skip Softmax

Skip Softmax 已集成到 NVIDIA TensorRT-LLM 中，可通过 LLM API 访问。用户可以配置稀疏注意力设置以根据特定需求优化性能。此功能在 NVIDIA 最新的数据中心 GPU 上得到支持，进一步加速注意力计算。

有关更多技术细节及开始使用 Skip Softmax，开发者可以参考 [NVIDIA 官方内容](https://developer.nvidia.com/blog/accelerating-long-context-inference-with-skip-softmax-in-nvidia-tensorrt-llm/)。