驯服指数：整数本机边缘推理的快速 Softmax 代理

Softmax 可能成为 Transformer 模型的多头注意力 (MHA) 块中的计算瓶颈，特别是在低精度推理下的小模型中，其中求幂和归一化会产生大量开销。因此，我们建议使用 Head-Calibrate Clipped-Linear Softmax (HCCS)，它是指数 softmax 函数的有界、单调代理，它使用最大中心注意力 logits 的剪辑线性映射。该近似产生稳定的概率分布，保持原始 logits 的排序并具有非负值。 HCCS 与以前的 softmax 代理不同，因为它包含一组轻量级校准参数，这些参数基于代表性数据集进行离线优化，并针对每个单独的注意力头进行校准，以保留各个头的统计属性。我们描述了针对 AMD Versal AI 引擎的高吞吐量场景的硬件驱动的 HCCS 实现。 AMD 目前针对该平台的参考实现依赖 bfloat16 算术或 LUT 来执行指数运算，这可能会限制平台的吞吐量，并且无法利用 AI 引擎的高吞吐量整数向量处理单元。相比之下，HCCS 提供了到 AI 引擎的 int8 乘法累加 (MAC) 单元的自然映射。据我们所知，这是 AMD AI 引擎的第一个 int8 优化的 softmax 代理，它显着超过了其他参考实现的速度性能，同时在量化感知再训练后在小型或大量量化的 MHA 工作负载上保持了有竞争力的任务准确性。