手写 PTX Tensor-Core GEMM 内核：NVIDIA L4 上的多精度研究

文章背景与核心概要

高性能 Tensor Core 内核通常依赖于底层 PTX 管道，该管道利用异步数据移动（cp.async）、warp 级矩阵加载（ldmatrix）以及矩阵乘累加操作（mma.sync）。然而，大多数应用程序是通过高级 WMMA C++ API 间接访问 Tensor Core 的。本文研究了一个实际的工程问题：用手写 PTX 代替 WMMA 在什么时候真正划算？

作者在 NVIDIA L4 GPU（Ada Lovelace 架构，SM89）上进行了受控的单 GPU 分析，对比了双缓冲 WMMA 基准测试与一系列手写 PTX GEMM（通用矩阵乘法）内核，涵盖 FP16、INT8 和 INT4 精度，问题规模从 \(N = 512\) 到 \(N = 8192\) 不等。研究表明，寄存器分配、数据打包开销以及内存合并访问模式在不同精度下对性能有着决定性影响，并为需要手写 PTX 的特定场景提供了清晰的性能参考。

摘要概要

High-performance Tensor Core kernels typically rely on a low-level PTX pipeline that utilizes asynchronous data movement (cp.async), warp-level matrix loads (ldmatrix), and matrix multiply-accumulate operations (mma.sync). However, most applications access Tensor Cores indirectly via the high-level WMMA C++ API.

This study investigates a practical engineering question: When does replacing WMMA with hand-written PTX actually pay off?

Conducting a controlled, single-GPU analysis on an NVIDIA L4 GPU (Ada Lovelace, SM89), the authors compare double-buffered WMMA baselines against a family of hand-written PTX GEMM (General Matrix Multiply) kernels across FP16, INT8, and INT4 arithmetic with square problem sizes ranging from \(N = 512\) to \(N = 8192\).

高性能 Tensor Core 内核通常依赖于底层 PTX 管道，该管道利用异步数据移动（cp.async）、warp 级矩阵加载（ldmatrix）以及矩阵乘累加操作（mma.sync）。然而，大多数应用程序是通过高级 WMMA C++ API 间接访问 Tensor Core 的。

High-performance Tensor Core kernels typically rely on a low-level PTX pipeline that utilizes asynchronous data movement (cp.async), warp-level matrix loads (ldmatrix), and matrix multiply-accumulate operations (mma.sync). However, most applications access Tensor Cores indirectly via the high-level WMMA C++ API.

本文研究了一个实际的工程问题：用手写 PTX 代替 WMMA 在什么时候真正划算？

This study investigates a practical engineering question: When does replacing WMMA with hand-written PTX actually pay off?

作者在 NVIDIA L4 GPU（Ada Lovelace，SM89）上进行了受控的单 GPU 分析，对比了双缓冲 WMMA 基准测试与一系列手写 PTX GEMM（通用矩阵乘法）内核，涵盖 FP16、INT8 和 INT4 精度，问题规模从 \(N = 512\) 到 \(N = 8192\) 不等。

Conducting a controlled, single-GPU analysis on an NVIDIA L4 GPU (Ada Lovelace, SM89), the authors compare double-buffered WMMA baselines against a family of hand-written PTX GEMM (General Matrix Multiply) kernels across FP16, INT8, and INT4 arithmetic with square problem sizes ranging from \(N = 512\) to \(N = 8192\).

关键发现与性能结果

Key Findings & Performance Results

FP16（半精度）： 手写 PTX 没有带来端到端的加速。指令级别的收益完全被操作数打包开销（operand-packing overhead）所抵消。
- FP16 (Half Precision): Hand-written PTX provides no end-to-end speedup. Instruction-level gains are entirely offset by operand-packing overhead.
INT8（8位整数）： PTX 内核实现了 1.4倍至1.8倍 的持续加速，这主要归功于更低的指令数和更优的全局内存合并访问（coalescing）。
- INT8 (8-bit Integer): PTX kernels achieve consistent speedups of 1.4x to 1.8x, driven primarily by lower instruction counts and superior global-memory coalescing.
INT4（4位整数）： PTX 内核实现了 2.9倍至4.3倍 的惊人加速，其中原生 mma.sync.m16n8k64.s4 绕过了标准 WMMA 路径所需的软件仿真序列。
- INT4 (4-bit Integer): PTX kernels achieve dramatic speedups of 2.9x to 4.3x, where native mma.sync.m16n8k64.s4 execution bypasses the software-emulated sequences required by the standard WMMA path.
量化峰值： 相对于 FP16 WMMA 基准，在 \(N = 8192\) 时，最优的量化内核分别达到了 34.4倍（INT8）和 98.7倍（INT4）的加速比。
- Quantized Peaks: Relative to the FP16 WMMA baseline, the best quantized kernels reach 34.4x (INT8) and 98.7x (INT4) speedups at \(N = 8192\).

GPU 性能指标的洞察

Insights on GPU Performance Metrics

从分析数据（通过 NVIDIA Nsight Compute 收集的全套指标）中得出的一个重要结论是：占有率（Occupancy）是吞吐量的糟糕预测指标。

A major takeaway from the profiling data (collected via NVIDIA Nsight Compute across its full metric set) is that occupancy is a poor predictor of throughput.

对于大型矩阵，性能表现主要跟随内存系统的行为——具体来说是全局加载合并（global-load coalescing）和 DRAM 活动周期（DRAM-active cycles）——其关联度远高于原始的 Tensor Core 利用率。这些见解成功地勾勒出了能够证明手写 PTX 带来的额外复杂性是合理的特定精度和运算区间。

For large matrices, performance tracks memory-system behavior—specifically global-load coalescing and DRAM-active cycles—far more closely than raw Tensor Core utilization. These insights successfully map out the specific precisions and operational regimes that justify the added complexity of hand-written PTX.