Tensor cores and matrix operations

What Are Tensor Cores?

How Matrix Multiplication Works (From Scratch)

Traditional Approach: Element-by-Element

For C = A × B where A is M×K, B is K×N, C is M×N:

$C[i][j] = \sum_{k=0}^{K-1} A[i][k] \cdot B[k][j]$

On CUDA cores: Each thread computes one C[i][j] by:

1. Loading row i of A (K elements)
2. Loading column j of B (K elements)
3. Performing K multiply-adds serially
4. Writing result to C[i][j]

Problem: For 1024×1024 matrices, that's 1024³ ≈ 1 billion FMAs (≈ 2 billion FLOPs). At 10 TFLOPS (10×10¹² FLOPS), that's about $2\times10^9 / 10^{13} \approx 0.2$ ms of pure compute — but memory stalls and poor reuse make the real CUDA-core time much worse.

Tensor Core Approach: Tiled Matrix Operations

Key insight: Break large matrix multiply into small tiles and use specialized hardware for each tile.

Complete Algorithm for Large Matrix Multiply

Problem: Multiply 1024×1024 matrices using 4×4 Tensor Cores.

Solution: Tiled multiply with accumulation

# Pseudocode for C = A × B using Tensor Cores
M, N, K = 1024, 1024, 1024
TILE = 4
 
for i in range(0, M, TILE):          # Tile rows of A
    for j in range(0, N, TILE):      # Tile columns of B
        D_tile = zeros(TILE, TILE)   # Accumulator
        for k in range(0, K, TILE):  # Tile inner dimension
            A_tile = A[i:i+TILE, k:k+TILE]
            B_tile = B[k:k+TILE, j:j+TILE]
            # SINGLE TENSOR CORE INSTRUCTION
            D_tile = mma_sync(A_tile, B_tile, D_tile)
        
        C[i:i+TILE, j:j+TILE] = D_tile

Why this works:

1. Tiling preserves correctness: Matrix multiply is associative, so we can sum partial products
2. Each iteration computes a 4×4 partial product: A[i:i+4, k:k+4] × B[k:k+4, j:j+4]
3. Accumulation in D_tile sums all K/4 partial products for final result
4. Outer loops cover all tiles of output matrix C

Performance calculation:

• Scalar FMAs: 1024³ ≈ 1.07 billion FMAs
• Number of Tensor Core ops: (1024/4)³ = 256³ ≈ 16.8 million
• Ratio: $1024^3 / 256^3 = 4^3 = 64$ → 64× instruction reduction compared to scalar (each Tensor Core op does 4×4×4 = 64 FMAs)

Precision Modes and Their Impact

Programming Tensor Cores

WMMA API (Warp-level Matrix Multiply Accumulate)

#include <mma.h>
using namespace nvcuda::wmma;
 
__global__ void matmul_tensor_core(
    half *A, half *B, float *C, int M, int N, int K) {
    // Each warp computes a 16×16 tile of C
    int warpM = (blockIdx.x * blockDim.x + threadIdx.x) / 32;
    int warpN = (blockIdx.y * blockDim.y + threadIdx.y);
    
    // Declare fragments (distributed across warp)
    fragment<matrix_a, 16, 16, 16, half, row_major> a_frag;
    fragment<matrix_b, 16, 16, 16, half, col_major> b_frag;
    fragment<accumulator, 16, 16, 16, float> c_frag;
    fill_fragment(c_frag, 0.0f);  // Initialize accumulator
    
    // Tile K dimension
    for (int i = 0; i < K; i += 16) {
        // Load tiles (each thread gets its portion)
        load_matrix_sync(a_frag, A + warpM*16*K + i, K);
        load_matrix_sync(b_frag, B + i*N + warpN*16, N);
        
        // TENSOR CORE OPERATION
        mma_sync(c_frag, a_frag, b_frag, c_frag);
    }
    
    // Store result
    store_matrix_sync(C + warpM*16*N + warpN*16, c_frag, N, mem_row_major);
}

Why each step?

• fragment: Data distributed across 32 threads of warp (each thread holds a piece of the tile)
• load_matrix_sync: Warp-level load, each thread fetches its portion from global memory
• mma_sync: The Tensor Core instruction — all 32 threads participate cooperatively
• store_matrix_sync: Warp-level store, accumulator fragment → global memory

cuBLAS and cuDNN Automatic Use

import torch
 
# PyTorch automatically uses Tensor Cores when:
# 1. Input dtype is float16, bfloat16, or TF32
# 2. Dimensions are multiples of 8
# 3. torch.backends.cuda.matmul.allow_tf32 = True (default on Ampere+)
 
A = torch.randn(1024, 1024, dtype=torch.float16, device='cuda')
B = torch.randn(1024, 1024, dtype=torch.float16, device='cuda')
 
# Uses Tensor Cores automatically
C = torch.matmul(A, B)

Why dimensions should be multiples of 8?

• Tensor Core API tiles are 16×16 (with K in multiples of 8/16 depending on precision)
• Padding to multiples ensures full tile coverage
• Partial tiles need fallback handling (performance cliff)

Common Mistakes and Fixes

Performance Characteristics

Memory Layout and Optimization

When to Use Tensor Cores vs CUDA Cores

Connections

• 6.2.1-GPU-vs-CPU-architecture — Why GPUs need specialized units
• 6.2.8-Memory-hierarchy-in-GPUs — How shared memory enables tiling
• 6.2.10-Warp-scheduling-and-execution — How warps map to Tensor Core fragments
• 8.3.4-Mixed-precision-training — Using Tensor Cores in deep learning
• 8.4.2-Quantization-techniques — INT8 and INT4 Tensor Core modes
• 9.1.5-Roofline-model — Analyzing Tensor Core performance limits
• 7.3.6-cuBLAS-and-cuDNN — Libraries that abstract Tensor Core usage

#flashcards/hardware

What is the theoretical speedup of A100 Tensor Cores over CUDA cores? :::

Concept Map

Hinglish (regional understanding)

Hinglish (regional understanding)