Prefer to read without ads? Become a member — from $10/month — and support the work. Already a member? Log in to read ad-free on this device.

17.5 Tensor Cores

The register-blocked kernel is close to the limit of the FP32 CUDA cores’ performance; to go even faster, we must enlist the Tensor Cores, which were built specifically for matrix multiplication. The Tensor Cores implement machine instructions that multiply matrix tiles in a single operation (dozens of multiply-add instructions), and the nvcuda::wmma API exposes it through fragments – opaque per-warp tiles that the entire warp loads, multiplies, and stores cooperatively.

Section 17.1 read C one element at a time – a row of A dotted with a column of B – but the product is equally a sum of outer products: over k, each term is the rank-1 product of column k of A and row k of B. A single Tensor Core instruction evaluates a rank-8 slab of that sum: a 16×16 output tile from a 16×8 fragment of A and an 8×16 fragment of B.

To keep things simple for illustrative purposes, this chapter only uses the TF32 feature of Tensor Cores (introduced in Ampere), which essentially coerces FP32 inputs to FP16 precision. Over time, the formats supported by Tensor Cores have steadily increased, with a trend toward smaller and smaller FP formats. The latest, FP4, can only represent numbers between +/-6.0! NVIDIA has been able to use these narrower FP formats to deliver higher peak FLOPS rates than otherwise would be possible.

For single precision, the natural entry point is TF32, a 19-bit input format with the 8-bit exponent of float (so the same dynamic range) but only a 10-bit mantissa. Inputs are rounded from float to TF32 on the way into the fragment; the multiply-accumulate happens in the Tensor Core; and the result accumulates in full FP32. The reduced input precision keeps float’s range (exponents are not changed), but discards thirteen bits of mantissa, so a TF32 GEMM is an approximate GEMM – but with precision enough for many deep-learning workloads.

for ( int kk = 0; kk < BK; kk += 8 ) {    wmma::fragment<wmma::matrix_a,16,16,8,wmma::precision::tf32,wmma::row_major> aF[TMw];    wmma::fragment<wmma::matrix_b,16,16,8,wmma::precision::tf32,wmma::row_major> bF[TNw];    for ( int i = 0; i < TMw; i++ ) {        wmma::load_matrix_sync( aF[i], &As[warpRow*(TMw*16)+i*16][kk], BK );        for ( int t = 0; t < aF[i].num_elements; t++ )            aF[i].x[t] = wmma::__float_to_tf32( aF[i].x[t] );      // round float -> TF32    }    for ( int j = 0; j < TNw; j++ ) {        wmma::load_matrix_sync( bF[j], &Bs[kk][warpCol*(TNw*16)+j*16], BN );        for ( int t = 0; t < bF[j].num_elements; t++ )            bF[j].x[t] = wmma::__float_to_tf32( bF[j].x[t] );    }    for ( int i = 0; i < TMw; i++ )        for ( int j = 0; j < TNw; j++ )            wmma::mma_sync( cF[i][j], aF[i], bF[j], cF[i][j] );     // accumulate in FP32}

Listing 17-5. Tensor-core inner loop, TF32 via WMMA (sgemm_wmma).

Wrapped in the same shared-memory staging as before, with a 128×64 block tile, the WMMA kernel reaches 8718 GFLOP/s – higher than cuBLAS’s FP32 number, using inputs the FP32 kernels could not. But there’s still performance to be gained: the Tensor Cores retire multiply-accumulates so quickly that the kernel can no longer keep them fed. With the arithmetic so cheap, the global-memory load of the next slab is once more the limiting factor. The kernel is memory-bound, and it stalls at each __syncthreads() waiting for data to arrive.

Figure 17-4. A warp-level MMA multiplies a 16×8 fragment of A by an 8×16 fragment of B into a 16×16 FP32 accumulator.