用于高效低位 MXFP 推理的对角平铺混合精度注意力

基于 Transformer 的 大语言模型 (LLM) 在各种现实任务中表现出了卓越的性能，但由于注意力的二次复杂度和高精度操作的内存带宽限制，其推理成本仍然过高。在这项工作中，我们利用下一代 GPU 架构的计算能力，提出了一种使用微尺度浮点 (MXFP) 数据格式的低位混合精度注意力内核。我们的对角平铺混合精度注意力（DMA）在平铺级别融合了两种低位计算，是使用 Triton 实现的精致融合内核，利用硬件级并行性和内存效率来实现快速高效的推理，而不会影响模型性能。对 NVIDIA B200 GPU 的广泛实证评估表明，我们的内核保持了生成质量，且性能下降可以忽略不计，同时通过内核融合实现了显着的加速。我们在 https://github.com/yifu-ding/MP-Sparse-Attn 发布了我们的代码。