利用 SplitK 工作分解加速用于 W4A16 量化推理的 Triton 融合核
分布式、并行与集群计算
2024-02-26 v2 人工智能
摘要
我们提出了一种用于 W4A16 量化推理的高效融合矩阵乘法核的实现方案，其中我们使用 SplitK 工作分解在融合核中执行反量化和 GEMM。我们的实现展示了在基础模型推理工作负载中常见的那种瘦长型矩阵-矩阵乘法的性能改进。具体而言，本文考察了瘦激活矩阵与方权重矩阵之间的矩阵乘法类型。我们的结果显示，在 A100 上平均速度提升 65%，在 H100 上平均速度提升 124%（峰值达 295%），适用于一系列矩阵维度，包括在 llama 风格模型中发现的那些维度，其中 m < n = k。
引用
@article{arxiv.2402.00025,
title = {Accelerating a Triton Fused Kernel for W4A16 Quantized Inference with SplitK work decomposition},
author = {Adnan Hoque and Less Wright and Chih-Chieh Yang and Mudhakar Srivatsa and Raghu Ganti},
journal= {arXiv preprint arXiv:2402.00025},
year = {2024}
}