AI驱动技术正日益融入日常生活，提升知识获取效率与生产力。大型语言模型（LLM）是这些应用的核心，但其内存密集型特性需依赖专用硬件加速器提供数十exaflops算力。本文通过优化内存使用应对计算挑战。

LLM的大部分内存和计算消耗于矩阵乘法操作中的权重。使用更窄的数据类型可减少内存占用，例如8位整数（U8/S8）存储权重相比单精度（F32）减少4倍内存，相比半精度（F16）或bfloat16（BF16）减少2倍。权重仅量化（weight-only quantization）技术将S8权重与F16输入结合，在精度可接受范围内提升效率，但需软件实现混合输入矩阵乘法（如F16输入乘8位整数）。硬件加速器（包括GPU）支持固定数据类型，因此混合输入需通过软件映射到硬件操作。

本文重点在NVIDIA Ampere架构上映射混合输入矩阵乘法，提出软件技术解决数据类型转换和布局一致性，以高效利用硬件支持的数据类型和布局。结果显示软件额外开销极小，性能接近硬件峰值。相关技术已开源至NVIDIA/CUTLASS库。

矩阵乘加操作

现代AI硬件加速器（如Google TPU和NVIDIA GPU）通过Tensor Core原生执行矩阵乘法。Ampere Tensor Core提供矩阵乘加（mma）操作，支持固定数据类型、形状和数据布局的输入矩阵。更大规模或不同数据类型的矩阵乘法需通过软件分块映射到硬件支持的操作。

Tensor Core mma操作定义两个输入矩阵（A和B）生成结果矩阵C，支持混合精度（输入与结果数据类型不同），但不支持混合输入（输入数据类型不同），需软件实现。

混合输入矩阵乘法的挑战

以F16输入乘U8权重（F16 * U8）为例，讨论挑战：

• 数据类型转换：mma操作要求输入数据类型相同，需将U8转换为F16以利用混合精度Tensor Core。权重数量庞大，转换操作需降低延迟。
• 布局一致性：mma操作要求输入矩阵在寄存器中的布局符合硬件规范。U8权重矩阵转换后需调整布局以匹配F16数据类型，需软件实现布局一致性。

软件策略应对挑战

典型数据类型转换涉及32位寄存器操作序列（如4xU8转2x(2xF16)），约需10条指令。布局一致性现有解决方案包括：

• 窄位宽共享内存加载：直接加载U8数据到寄存器，避免重排操作，但未充分利用共享内存带宽。
• 全局内存预处理：在全局内存重排数据，实现宽位宽加载，最大化带宽利用，但需额外硬件特定预处理步骤（如NVIDIA/FasterTransformer采用此法）。

优化软件策略

我们实现FastNumericArrayConvertor和FragmentShuffler进一步优化：

• FastNumericArrayConvertor：在32位寄存器上操作4xU8，避免解包单个值，使用低成本算术操作减少指令数并加速转换。U8转F16约用6条指令，减少1.6倍。
• FragmentShuffler：通过重排数据实现布局一致性，允许宽位宽加载操作，提高共享内存带宽利用并减少操作总数。利用ldmatrix指令从共享内存加载数据到寄存器，再通过shuffle操作（shfl.sync）实现布局一致性，避免全局内存预处理或窄位宽加载。

本工作最大贡献是通过寄存器重排实现布局一致性，覆盖U8-to-F16、S8-to-F16、U8-to-BF16和S8-to-BF16转换。

性能结果

在NVIDIA A100 SXM芯片上测量八种混合输入变体（蓝红色）和两种混合精度（绿色）的性能（FLOPS）。尽管前八种需额外操作（数据类型转换和布局一致性），但性能接近或与混合精度相当。

致谢

感谢Quentin Colombet、Jacques Pienaar等团队成员的技术贡献，以及NVIDIA合作伙伴Rawn Henry、Pradeep Ramani等的支持。