Diagonal-Tiled Mixed-Precision Attention for Efficient Low-Bit MXFP Inference

📄 arXiv: 2604.03950 📥 PDF

作者: Yifu Ding, Xinhao Zhang, Jinyang Guo

分类: cs.LG, cs.AI

发布日期: 2026-04-07

💡 一句话要点

提出对角分块混合精度注意力机制，加速低比特MXFP大模型推理。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 大语言模型 低比特量化 混合精度计算 注意力机制 GPU推理 MXFP Triton 核融合

📋 核心要点

  1. 大语言模型推理成本高昂，主要瓶颈在于注意力机制的计算复杂度和高精度运算的内存带宽限制。
  2. 提出对角分块混合精度注意力(DMA)机制，利用低比特MXFP格式，在分块级别上进行低比特计算。
  3. 在NVIDIA B200 GPU上实验表明，DMA在保持生成质量的同时，通过核融合实现了显著的推理加速。

📝 摘要（中文）

基于Transformer的大语言模型(LLMs)在各种实际任务中表现出色，但由于注意力机制的二次复杂度以及高精度运算的内存带宽限制，其推理成本仍然过高。本文提出了一种使用微缩放浮点(MXFP)数据格式的低比特混合精度注意力核，充分利用下一代GPU架构的计算能力。我们的对角分块混合精度注意力(DMA)在分块级别上结合了两种低比特计算，并使用Triton实现了一个精巧的融合核，利用硬件级并行性和内存效率来实现快速高效的推理，同时不影响模型性能。在NVIDIA B200 GPU上的大量实验评估表明，我们的核保持了生成质量，且性能几乎没有下降，同时通过核融合实现了显著的加速。代码已开源。

🔬 方法详解

问题定义：现有大语言模型推理面临高昂的计算成本，尤其是在注意力机制中，其计算复杂度是序列长度的平方。此外，高精度（如FP32或FP16）运算对内存带宽的需求很高，进一步限制了推理速度。因此，如何在保证模型性能的前提下，降低注意力机制的计算复杂度和内存带宽需求，是本文要解决的核心问题。

核心思路：本文的核心思路是利用低比特混合精度计算来降低计算复杂度和内存带宽需求。具体而言，采用微缩放浮点(MXFP)数据格式，并在注意力计算中引入对角分块策略，在分块级别上进行低比特计算。通过这种方式，可以在保证模型性能的同时，显著降低计算量和内存访问量。

技术框架：DMA (Diagonal-Tiled Mixed-Precision Attention) 的整体框架包括以下几个主要步骤：1) 将输入 Query, Key, Value 矩阵进行分块；2) 在每个分块内，使用低比特 MXFP 格式进行注意力计算；3) 对角分块策略，允许在不同的分块上使用不同的精度，从而实现混合精度计算；4) 使用 Triton 语言实现高度优化的融合核，充分利用硬件级的并行性和内存效率。

关键创新：本文最重要的技术创新点在于提出了对角分块混合精度注意力(DMA)机制。与传统的全精度或固定精度量化方法不同，DMA 允许在不同的分块上使用不同的精度，从而实现更灵活的精度控制。此外，DMA 采用对角分块策略，可以更好地适应不同序列长度的输入，并提高计算效率。

关键设计：DMA 的关键设计包括：1) MXFP 数据格式的选择，MXFP 是一种专门为深度学习推理设计的低比特浮点格式，可以在保证模型性能的同时，显著降低计算量和内存带宽需求；2) 对角分块策略的设计，对角分块可以更好地利用 GPU 的并行计算能力，并减少内存访问的冲突；3) Triton 融合核的实现，Triton 是一种专门为 GPU 编程设计的语言，可以实现高度优化的核函数，充分利用硬件资源。

🖼️ 关键图片

fig_0
fig_1

📊 实验亮点

在NVIDIA B200 GPU上的实验结果表明，DMA在保持生成质量几乎没有下降的情况下，实现了显著的推理加速。具体而言，DMA相比于传统的全精度注意力机制，可以实现高达X倍的加速（具体数值需要在论文中查找），同时模型性能的下降可以忽略不计。这些结果表明，DMA是一种高效且实用的低比特混合精度注意力机制。

🎯 应用场景

该研究成果可广泛应用于各种需要高效大语言模型推理的场景，例如：移动设备上的本地推理、边缘计算设备上的实时翻译、以及云计算平台上的大规模模型部署。通过降低推理成本，可以促进大语言模型在更多实际应用中的普及，并加速人工智能技术的发展。

📄 摘要（原文）

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 atthis https URL.