手写PTX代码在低精度计算中显著提速,适合量化推理场景。
Hand-Written PTX Tensor-Core GEMM Kernels: A Multi-Precision Study on NVIDIA L4

- 直接用PTX重写Tensor Core核函数,绕过高级API开销。
- INT8下提速1.4至1.8倍,INT4达2.9至4.3倍,最大超FP16基线98.7倍。
- 性能主要受内存访问效率影响,而非核心利用率,适合大矩阵量化任务。
高性能Tensor Core内核依赖低级PTX流水线,包括异步数据移动(cp.async)、波前级矩阵加载(ldmatrix)和矩阵乘累加(mma.sync)。但多数应用通过WMMA C++ API间接调用。本文在NVIDIA L4 GPU(Ada, SM89)上开展单卡对照实验,比较双缓冲的WMMA基线与一系列手写PTX GEMM核函数在FP16、INT8、INT4精度及512至8192平方问题规模下的表现。所有核函数均使用Nsight Compute全量指标分析,速度提升以同精度WMMA基线为基准。结果显示:FP16无端到端加速,因指令级优势被操作数打包开销抵消;而INT8下稳定提速1.4x–1.8x,源于更低指令数与更优全局内存合并;INT4下提速2.9x–4.3x,得益于原生mma.sync.m16n8k64.s4执行避免了软件模拟路径。相较FP16 WMMA基线,最优量化核在N=8192时分别达到34.4x(INT8)和98.7x(INT4)加速。实验表明,占用率无法预测吞吐,大矩阵性能更依赖内存系统行为——尤其是全局加载合并与DRAM活跃周期,而非Tensor Core利用率。该结果明确了手写PTX在特定精度与场景下的合理性。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。