用低比特混合精度计算加速大模型注意力，兼顾速度与质量。
Diagonal-Tiled Mixed-Precision Attention for Efficient Low-Bit MXFP Inference
- 在分块层级引入两种低比特计算，结合Triton实现融合内核。
- 在NVIDIA B200上实现显著提速，生成质量几乎无损。
- 适合追求高效推理的LLM部署场景，尤其看重内存效率者。
基于Transformer的大语言模型在众多实际任务中表现卓越，但其推理成本因注意力机制的二次复杂度以及高精度运算对内存带宽的限制而居高不下。本文提出一种基于微缩浮点（MXFP）数据格式的低比特混合精度注意力内核，利用下一代GPU架构的计算能力。所提出的对角分块混合精度注意力（DMA）在分块层级引入两种低比特计算，采用Triton实现精细融合内核，充分利用硬件级并行与内存效率，实现快速高效的推理且不牺牲模型性能。在NVIDIA B200 GPU上的大量实证评估表明，该内核在保持生成质量近乎无损的同时，通过内核融合实现了显著加速。代码已开源：https://github.com/yifu-ding/MP-Sparse-Attn。
原文摘要 · Abstract (English)
Transformer-based large language models (LLMs) have demonstrated remarkable performance across a wide range of real-world tasks, but their inference cost remains prohibitively high due to the quadratic complexity of attention and the memory bandwidth limitations of high-precision operations. In this work, we present a low-bit mixed-precision attention kernel using the microscaling floating-point (MXFP) data format, utilizing the computing capability on next-generation GPU architectures. Our Diagonal-Tiled Mixed-Precision Attention (DMA) incorporates two kinds of low-bit computation at the tiling-level, and is a delicate fused kernel implemented using Triton, exploiting hardware-level parallelism and memory efficiency to enable fast and efficient inference without compromising model performance. Extensive empirical evaluations on NVIDIA B200 GPUs show that our kernel maintains generation quality with negligible degradation, and meanwhile achieves significant speedup by kernel fusion. We release our code at https://github.com/yifu-ding/MP-Sparse-Attn.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。