Tensor cores and matrix operations

6.2.12 · Hardware › GPU Architecture

Tensor Cores Kya Hain?

Matrix Multiplication Kaise Kaam Karta Hai (Scratch Se)

Traditional Approach: Element-by-Element

C = A × B ke liye jahan A, M×K hai, B, K×N hai, C, M×N hai:

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

CUDA cores par: Har thread ek C[i][j] compute karta hai:

1. A ki row i load karna (K elements)
2. B ka column j load karna (K elements)
3. K multiply-adds serially perform karna
4. Result C[i][j] mein likhna

Problem: 1024×1024 matrices ke liye, yeh 1024³ ≈ 1 billion FMAs hain (≈ 2 billion FLOPs). 10 TFLOPS (10×10¹² FLOPS) par, yeh pure compute ke liye lagbhag $2\times 10^9/10^{13}\approx 0.2$ ms hai — lekin memory stalls aur poor reuse real CUDA-core time ko kaafi bura bana dete hain.

Tensor Core Approach: Tiled Matrix Operations

Key insight: Badi matrix multiply ko chhoti tiles mein todna aur har tile ke liye specialized hardware use karna.

Badi Matrix Multiply ke liye Complete Algorithm

Problem: 4×4 Tensor Cores use karke 1024×1024 matrices multiply karo.

Solution: Accumulation ke saath tiled multiply

# 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

Yeh kyun kaam karta hai:

1. Tiling correctness preserve karta hai: Matrix multiply associative hai, isliye hum partial products sum kar sakte hain
2. Har iteration ek 4×4 partial product compute karta hai: A[i:i+4, k:k+4] × B[k:k+4, j:j+4]
3. Accumulation D_tile mein final result ke liye saare K/4 partial products sum karta hai
4. Outer loops output matrix C ke saare tiles cover karte hain

Performance calculation:

• Scalar FMAs: 1024³ ≈ 1.07 billion FMAs
• Tensor Core ops ki number: (1024/4)³ = 256³ ≈ 16.8 million
• Ratio: $1024^3/256^3=4^3=64$ → scalar ke comparison mein 64× instruction reduction (har Tensor Core op 4×4×4 = 64 FMAs karta hai)

Precision Modes aur Unka Impact

Tensor Cores Ko Program Karna

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);
}

Har step kyun?

• fragment: Data warp ke 32 threads mein distributed hai (har thread tile ka ek piece hold karta hai)
• load_matrix_sync: Warp-level load, har thread global memory se apna portion fetch karta hai
• mma_sync: The Tensor Core instruction — saare 32 threads cooperatively participate karte hain
• store_matrix_sync: Warp-level store, accumulator fragment → global memory

cuBLAS aur 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)

Dimensions 8 ke multiples kyun hone chahiye?

• Tensor Core API tiles 16×16 hoti hain (precision ke hisaab se K, 8/16 ke multiples mein)
• Multiples tak padding full tile coverage ensure karti hai
• Partial tiles ko fallback handling chahiye (performance cliff)

Common Mistakes aur Fixes

Performance Characteristics

Memory Layout aur Optimization

Tensor Cores vs CUDA Cores Kab Use Karein

Connections

• 6.2.1-GPU-vs-CPU-architecture — GPUs ko specialized units kyun chahiye
• 6.2.8-Memory-hierarchy-in-GPUs — Shared memory tiling ko kaise enable karta hai
• 6.2.10-Warp-scheduling-and-execution — Warps Tensor Core fragments se kaise map hote hain
• 8.3.4-Mixed-precision-training — Deep learning mein Tensor Cores use karna
• 8.4.2-Quantization-techniques — INT8 aur INT4 Tensor Core modes
• 9.1.5-Roofline-model — Tensor Core performance limits analyze karna
• 7.3.6-cuBLAS-and-cuDNN — Libraries jo Tensor Core usage abstract karti hain

#flashcards/hardware

A100 Tensor Cores ka CUDA cores par theoretical speedup kya hai? :::

Concept Map