【翻译】Accelerating Llama3 FP8 Inference with Triton Kernels

文章预览

翻译：https://pytorch.org/blog/accelerating-llama3/ Accelerating Llama3 FP8 Inference with Triton Kernels by Adnan Hoque, Less Wright, Chih Chieh Yang 1.0 总结 我们提出了一种优化的Triton FP8 GEMM（通用矩阵乘法）kernel TK-GEMM，它利用了SplitK并行化。对于小batch size推理，在NVIDIA H100 GPU上针对Llama3-70B，TK-GEMM相比基础Triton矩阵乘法实现可提供高达 1.94 倍的性能提升，比cuBLAS FP8快 1.87 倍，比 cuBLAS FP16 快1.71倍。 图1. TK-GEMM相对于PyTorch（调用cuBLAS）在Llama3-70B注意力层矩阵形状下的加速比（N=K=8192） 在这篇博客中，我们将介绍如何使用Triton为FP8推理设计一个高效的kernel，并针对Llama3-70B推理进行调优。我们将讨论FP8（8位浮点数），这是Hopper代GPU（SM90）支持的新数据类型，Triton支持的SM90的关键特性，以及我们如何修改并行化以便能够为内存受限（推理）问题规模最大化内存吞吐量。 我们还专 ………………………………