arXivDaily arXiv每日学术速递 周一至周五更新
arXiv周末暂无论文更新,休息一下吧,周末愉快~~
arXiv 2608.10103cs.DCcs.AI

手写PTX张量核心GEMM内核:NVIDIA L4上的多精度研究

Hand-Written PTX Tensor-Core GEMM Kernels: A Multi-Precision Study on NVIDIA L4

Matt J. Borowski, Blazej Osinski

AI总结:

本文针对NVIDIA L4 GPU,对比手写PTX GEMM内核与WMMA基线,发现其在INT8、INT4精度下有显著加速,FP16下无增益,明确了手写PTX适用的精度与场景。

AI中文摘要:

高性能张量核心内核依赖于由异步数据移动(基于https URL)、使用ldmatrix的warp级矩阵加载,以及使用mma.sync.m16n8k64.s4等的矩阵乘累加操作构建的低级PTX流水线。然而,大多数应用代码通过WMMA C++ API间接访问张量核心。本文提出一个聚焦且实用的问题:用手写PTX替代WMMA何时真正划算?为回答该问题,我们在NVIDIA L4 GPU(Ada架构,SM89)上开展受控单GPU研究,在FP16、INT8、INT4算术精度及N=512到N=8192的方阵问题规模下,将双缓冲WMMA基线与一系列手写PTX GEMM内核进行对比。每个内核均用Nsight Compute在完整指标集下进行性能分析,PTX加速比是相对于对应同精度WMMA基线的结果。手写PTX在FP16下无端到端加速,因其指令级增益被操作数打包开销抵消;相反,PTX内核在INT8下实现1.4倍至1.8倍的一致加速,主要源于更少的指令数和更优的全局内存合并,在INT4下实现2.9倍至4.3倍的加速,其中原生mma.sync.m16n8k64.s4执行避免了WMMA路径使用的软件模拟序列。相对于FP16 WMMA基线,最优量化内核在N=8192时达到INT8下34.4倍、INT4下98.7倍的加速。在这些实验中, occupancy(占用率)是吞吐量的不良预测因子;对于大矩阵,性能更多取决于内存系统行为——尤其是全局加载合并和DRAM活跃周期——而非张量核心利用率。这些结果明确了手写PTX的额外复杂度在哪些精度和操作场景下是合理的。

英文摘要:

High-performance Tensor Core kernels rely on a low-level PTX pipeline built from asynchronous data movement with cp.async, warp-level matrix loads with ldmatrix, and matrix multiply-accumulate operations with mma.sync. However, most application code accesses Tensor Cores indirectly through the WMMA C++ API. This paper asks a focused, practical question: when does replacing WMMA with hand-written PTX actually pay off? To answer this question, we conduct a controlled, single-GPU study on an NVIDIA L4 GPU (Ada, SM89), comparing double-buffered WMMA baselines with a family of hand-written PTX GEMM kernels across FP16, INT8, and INT4 arithmetic and square problem sizes from $N=512$ to $N=8192$. Every kernel is profiled with Nsight Compute across the full metric set, and PTX speedups are reported relative to the corresponding same-precision WMMA baseline. Hand-written PTX provides no end-to-end speedup for FP16, because its instruction-level gains are offset by operand-packing overhead. In contrast, the PTX kernels achieve consistent speedups of 1.4x-1.8x for INT8, driven primarily by lower instruction counts and better global-memory coalescing, and 2.9x-4.3x for INT4, where native mma.sync.m16n8k64.s4 execution avoids the software-emulated sequence used by the WMMA path. Relative to the FP16 WMMA baseline, the best quantized kernels reach 34.4x (INT8) and 98.7x (INT4) at $N=8192$. Across these experiments, occupancy is a poor predictor of throughput. For large matrices, performance instead tracks memory-system behavior -- particularly global-load coalescing and DRAM-active cycles -- more closely than Tensor Core utilization. These results identify the precisions and operating regimes in which the additional complexity of hand-written PTX is justified.

补充信息

↑