对角分块混合精度注意力用于高效低位MXFP推理
机器学习
2026-04-07 v1 人工智能
摘要
基于Transformer的大语言模型（LLMs）在广泛的任务中展现了显著性能，但注意力机制的二次复杂度及高精度操作的内存带宽限制使其推理成本居高不下。本工作中，我们提出了一种使用微缩放浮点（MXFP）数据格式的低位混合精度注意力核，利用下一代GPU架构的计算能力。我们的对角分块混合精度注意力（DMA）在分块级别引入了两种低位计算，是一个使用Triton实现的精细融合核，利用硬件级并行性和内存效率实现快速高效推理，同时不损害模型性能。在NVIDIA B200 GPU上的广泛实证评估表明，我们的核在保持生成质量的同时几乎无退化，并通过核融合实现了显著加速。我们将代码发布在 https://github.com/yifu-ding/MP-Sparse-Attn。
引用
@article{arxiv.2604.03950,
title = {Diagonal-Tiled Mixed-Precision Attention for Efficient Low-Bit MXFP Inference},
author = {Yifu Ding and Xinhao Zhang and Jinyang Guo},
journal= {arXiv preprint arXiv:2604.03950},
year = {2026}
}
备注
CVPR Workshop EDGE 2026