为 AMD MI300 创建自定义内核

Hugging Face 与 AMD 开发了一套针对 AMD MI300X 的开源优化内核，以提升在 VLLM 中使用 FP8 推理 Llama 3.1 405B 的服务性能。通过实现三种特定的自定义内核——融合残差连接/RMS 归一化/FP8 转换内核、融合 SwiGLU 激活/FP8 转换内核以及 Skinny GEMM 内核——团队在解码阶段实现了显著的延迟降低（测量条件为输入大小 1，输出大小 128）。

自定义内核实现与性能提升

融合 RMS 归一化内核

RMS 归一化内核通过将残差连接、逐行均方根（RMS）归一化以及 FP8 量化融合为单一操作，优化了解码块的开头。

技术优化：

• 向量化内存访问： 该内核使用 128 位宽的加载指令，每次获取 8 个 FP16 元素，确保内存访问是合并且连续的，以最大化 warp 效率。
• 共享内存（SMEM）利用： 为避免重复访问 VRAM（全局内存），隐藏状态 $x$ 的修改版本存储在共享内存中。对于 Llama 405B，维度 $d=16384$ 能够适配每个计算单元可用的 64KB 共享内存。
• 块级归约： 内核为每行分配一个线程块，并使用共享内存同步线程，以完成 RMS 归一化所需的求和。

结果： “向量化 + SMEM”实现相较于标准 PyTorch 和 VLLM 现有实现都有显著提升，在各种批量大小下均提供了显著的加速。

融合 SwiGLU 内核

SwiGLU 内核将激活函数与随后对 MLP 块的 “Gate / Up” 投影的 FP8 量化融合在一起。

技术优化：

• 打包指令： 该内核利用 MI300X 的打包指令进行 FP16 加法和乘法，以提升每条指令的工作量。
• 快速数学近似： 为降低延迟，内核将标准的 exp 指令替换为更快的 exp2 指令，通过将输入乘以 $\log(2)$ 实现，精度损失可忽略不计。
• 打包 FP32 到 FP8 转换： 由于 MI300X 仅支持从 FP32 到 FP8 的转换，内核利用打包转换指令提升性能。

结果： 自定义的 SwiGLU 内核平均比 PyTorch 快超过 14 倍，且比 VLLM 内核快 27% 到 100% 不等。

Skinny GEMM 内核

标准库中的通用矩阵乘法（GEMM）内核对 “瘦” 矩阵（即行数极少的矩阵，解码时低批量大小常见）往往效率低下，因为它们的平铺机会受限，导致 GPU 利用率低。

技术优化：

• Split‑K 算法： 内核沿共享的 K 轴将 GEMM 拆分为多个子 GEMM 并行执行。通过分配工作负载，增加活跃计算单元（CUs）的数量，减少每个 CU 在 K 轴上的耗时。
• 稀疏技巧用于去除填充： 当行数少于最小稠密张量核心指令大小（例如 16）时，为避免浪费填充，内核使用 4:2 结构化稀疏指令。将稠密的 8 行矩阵映射为 16 行稀疏矩阵，从而能够使用 16x16x64 稀疏指令，其深度是最小稠密指令的两倍。
• Warp 专化与异步执行： 为应对低算术强度，内核将 warp 分为 “生产者”（专门从 VRAM 加载数据到共享内存）和 “消费者”（专门进行计算）。共享内存中的队列异步协调这些 warp，确保消费者在等待慢速 VRAM 加载时不会空闲。

结果： Skinny GEMM 内核在行数较少（M = 1、8、16）时相较于 PyTorch 显示出显著加速，尤其在 QKV 和 Gate/Up 投影上，但随着批量大小增至 32，提升幅度会减小。

实现与可用性

所有已开发的内核均可在 hf-rocm-kernels GitHub 仓库中获取，仓库包含源代码、Python 绑定、基准测试脚本以及测试套件。这些内核可独立使用，也可集成到 VLLM 中。要复现结果，Hugging Face 建议使用开发期间使用的特定 ROCm 6.3.1 容器。