× Install ThecoreGrid App
Tap below and select "Add to Home Screen" for full-screen experience.
B2B Engineering Insights & Architectural Teardowns

PTX Tensor Core GEMM: When Hand-Written Code Pays Off

Hand-written PTX Tensor-Core GEMM kernels demonstrate that the benefits of low-level code depend not on “closeness to the hardware,” but on what it removes from the execution path. This is particularly evident on NVIDIA L4 for INT8 and INT4.

This paper compares hand-written PTX Tensor-Core GEMM kernels with the baseline WMMA path on a single GPU, NVIDIA L4 on Ada SM89. The importance lies not in the fact of acceleration itself, but in the reason why it appears or disappears. The authors examine FP16, INT8, and INT4 across sizes from N = 512 to N = 8192, looking not only at wall time but also at memory behavior, caches, and execution counters.

The problem arises where WMMA ceases to be a convenient abstraction and becomes a limitation. The high-level API fixes the shape of fragments, load layout, and accumulator layout. This simplifies the code but leaves less room for optimization. PTX provides direct control over cp.async, ldmatrix, and mma.sync, but this comes at the cost of complexity, manual management of tiling, and the risk of making the kernel worse rather than better.

The authors address a practical question: when does replacing WMMA with PTX truly yield end-to-end gains? Their approach is pragmatic. They maintain the same block and warp tiling scheme, compare each PTX kernel only with the WMMA baseline of the same precision, and change one parameter at a time. This is a good engineering move because it allows linking the result to a specific cause rather than a set of randomly improved details.

For FP16, the conclusion was straightforward: hand-written PTX does not outperform compiler-optimized WMMA baseline. At small and medium sizes, the kernel is in a compute-bound mode, but even there, the gains from lower-level instructions are consumed by overhead on packing operands. At larger sizes, the system hits memory limits. When the working set exceeds L2, the differences between implementations nearly vanish. In other words, if the bottleneck is already in DRAM, optimizing the instruction stream does not change the overall picture.

INT8 shows a more interesting compromise. The best kernel, int8_ptx_mma_k32, is consistently faster than WMMA across all sizes. The reason lies not in the magic of Tensor Cores, but in a more compact instruction count and nearly perfect global-load coalescing. The authors separately demonstrate that occupancy does not explain throughput here: the highest occupancy is found in the WMMA kernel, but it is not the fastest. This is an important observation for SRE and performance engineers. High occupancy can be a good signal, but it is not always a target metric.

INT4 yields the most significant effect. Here, the gains are not just due to better scheduling, but because the WMMA path for s4 is effectively software-expanded. That is, it carries extra instructions and lane-dependent divergence. The native PTX kernel uses mma.sync.m16n8k64.s4 directly, removing the layer of emulation. As a result, speedups range from 2.9× to 4.3× compared to INT4 WMMA, and for FP16 WMMA, the best quantized kernels reach up to 98.7× at N = 8192. However, the picture is not linear: for small N, a three-stage pipeline is better, while for large N, the k64 variant wins, which better preserves L1 locality.

A separate interest lies in the ablation around the INT4 k64 family. The authors examine split loader, cache-eviction policy, and layout of the B-operand. The conclusion is neat: the chosen configuration is already at a local optimum. Changes in policy or layout do not yield improvements, and transposed B sharply worsens coalescing and cuts usable bandwidth. This is a useful result precisely because it dispels the illusion of “one more simple tuning.” Sometimes the optimal point has already been found, and the next step merely redistributes pressure between L1, L2, and issue slots.

The main takeaway of the article aligns well with engineering practice. PTX is justified when it removes overheads that WMMA cannot hide: emulation, extra instructions, poor coalescing. For FP16, this does not hold. For INT8, it provides a moderate but stable gain. For INT4, it offers a structural advantage over the native MMA path. Meanwhile, occupancy remains a diagnostic signal rather than a goal. For inference on quantized LLM on a single L4-class GPU, the authors recommend first examining memory-system behavior, and then Tensor Core utilization.


Information source

arXiv is the largest open preprint repository (since 1991, under the auspices of Cornell), where researchers quickly post working versions of papers; the materials are publicly accessible but do not undergo full peer review, so results should be considered preliminary and, where possible, checked against updated versions or peer‑reviewed journals. arxiv.org

View the original research PDF

×

🚀 Deploy the Blocks

Controls: ← → to move, ↑ to rotate, ↓ to drop.
Mobile: use buttons below.