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

This paper presents 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, for square problem sizes from N=512 to N=8192. The study asks a focused, practical question: when does replacing WMMA with hand-written PTX actually pay off?

FP16 (Run 1): Hand-written PTX provides no end-to-end speedup for FP16. Across all sizes, no hand-written PTX variant is faster than the compiler-optimized WMMA baseline in wall time; measured ratios stay within roughly ±5%. The reason is that in the compute-bound regime (N ≤ 4096) all kernels sustain 1,280–1,500 GFLOPS, and although the manual-pack variant reaches the highest SM throughput (57%), it also carries the most packing overhead, costing up to +22% at N=512. Once the 537 MB combined working set at N=8192 overflows L2, every kernel stalls on DRAM and all differences shrink below 1%.

INT8 (Run 2): The PTX kernels achieve consistent speedups of 1.4×–1.8×, driven primarily by lower instruction counts and better global-memory coalescing. The fastest kernel, int8 ptx mma k32, is 23% faster than int8 wmma at N=512 to 43% faster at N=8192. Its advantage comes from decomposing each K=32 tile step into two tightly unrolled m16n8k16 MMA calls, executing 25–43% fewer instructions than WMMA, and achieving near-perfect global-load coalescing with only 0.4% wasted sectors at N=8192, versus roughly 50% for every other kernel.

INT4 (Run 3): The largest gains appear in INT4, with speedups of 2.9×–4.3×. This is because the WMMA INT4 path (wmma::experimental::precision::s4) is software-expanded, inflating instruction count and introducing heavy lane-dependent divergence. The native PTX kernels call mma.sync.m16n8k64.s4 directly and avoid that expansion entirely. At N=8192, int4 wmma runs with only "16.89 active threads per warp and 4.6 million divergent branches, while the winning kernels sustain the full 32 active threads per warp with zero divergent branches and execute roughly 6–7× fewer instructions per scheduler."

Relative to FP16 WMMA baseline: The best quantized kernels reach 34.4× (INT8) and 98.7× (INT4) at N=8192.

The paper finds that 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. At N=8192, the Average DRAM Active Cycles metric is an almost exact fingerprint of wall time—k32 logs 1.57 B DRAM-active cycles versus WMMA's 2.80 B, a −43.9% reduction that matches its −43% wall-clock speedup.

The ablation over the best INT4 kernel family (loader split, cache-eviction policy, and operand layout) shows that the chosen configuration sits in a genuine local-optimum basin. The transposed-B variant collapses: with only 2.2 of every 32 bytes per global-load sector used, eligible warps per scheduler fall from 0.57 to 0.20 and warp cycles per instruction more than double, making it up to 3.1× slower at large N.

The paper's organizing principle: PTX is worth its complexity exactly when it removes instruction overhead that the WMMA path cannot avoid. For FP16, the compiler's WMMA lowering is already efficient and the workload becomes bandwidth-bound at scale, so PTX never wins. For INT8, the win is moderate and mechanistic. For INT4, the win is large and structural because WMMA emulates the s4 operation in software while PTX issues a single native MMA instruction. The practical takeaways for practitioners are: prefer native INT4/INT8 MMA over WMMA emulation; tune the K-tile and loader split for coalescing before reaching for deeper pipelines; and treat occupancy as a diagnostic, not a target.

Improvement 1: Adaptive Kernel Selection for Mixed-Precision GEMM

The improved AI system can automatically choose between WMMA and hand-written PTX kernels based on data type and problem size. For FP16, it defaults to WMMA (no speedup from PTX). For INT8, it selects int8 ptx mma k32 (1.4×–1.8× faster). For INT4, it uses native PTX mma.sync.m16n8k64.s4 (2.9×–4.3× faster). The system can also switch to memory-optimized kernels when matrix size exceeds L2 capacity (e.g., N > 4096), where DRAM coalescing dominates.

Improvement 2: Instruction-Count-Aware Compiler Optimization

The AI system can use the paper’s finding that PTX wins only when it removes instruction overhead. It can profile candidate kernels for instruction count per MMA tile and divergent branches. For INT4, it will reject WMMA’s software-expanded s4 path (16.89 active threads/warp, 4.6M divergent branches) and prefer native PTX (32 active threads, zero divergence, 6–7× fewer instructions). For FP16, it will skip PTX entirely since WMMA already achieves near-peak throughput.

Improvement 3: Coalescing-First Memory Scheduler

The improved AI system can prioritize global-load coalescing over occupancy when scheduling GEMM kernels. It uses the paper’s metric: wasted sectors per load. For INT8, it will adopt the k32 loader split that achieves 0.4% wasted sectors (vs. 50% for others). For INT4, it will use the transposed-A layout (not transposed-B, which collapses to 2.2 bytes/sector and 3.1× slowdown). The system can predict DRAM-active cycles as a proxy for wall time and tune tile sizes accordingly.

Improvement 4: Ablation-Based Kernel Configuration Tuner

The AI system can automatically run a mini-ablation (loader split, cache-eviction policy, operand layout) before deployment, as done in Run 4. It will detect local optima basins—e.g., for INT4, it will verify that the chosen config maintains ≥0.57 eligible warps/scheduler and <2× warp cycles per instruction. If a config violates these thresholds (like transposed-B), it will revert to the safer baseline.

Improved System Capability Summary

The enhanced AI system can:

• Deliver up to 4.3× speedup for INT4 GEMM and 1.8× for INT8 on NVIDIA L4 GPUs, with no regression on FP16.

• Automatically adapt to memory-bound regimes (N > 4096) by switching to coalescing-optimized kernels, reducing DRAM-active cycles by 44%.

• Reduce kernel development time by using instruction-count and coalescing heuristics to avoid trial-and-error PTX tuning.

• Provide a diagnostic dashboard showing occupancy vs. DRAM-active cycles vs. wall time, enabling practitioners to treat occupancy as a diagnostic, not a target.