⚡ Custom HGEMM: Half-Precision GEMM with Tensor Cores

A production-grade, highly-optimized Half-Precision General Matrix Multiplication (HGEMM) CUDA kernel built from scratch. This project demonstrates mastery over NVIDIA GPU architectures, shared memory management, WMMA, CuTe, asynchronous pipelining, and bank conflict resolution.

---

🏗️ Architecture & Progression

We systematically implement the GEMM kernel, optimizing iteratively to approach theoretical peak performance.

graph TD;
    A[Phase 1: Naive Global Memory] -->|Add Shared Memory| B(Phase 2: SMEM Tiling)
    B -->|Use Tensor Cores| C(Phase 3: WMMA API)
    C -->|Hide Latency| D(Phase 4: Async Pipeline)
    D -->|Eliminate Bank Conflicts| E(Phase 5: SMEM Swizzling)
    E -->|Rewrite with Abstractions| F(Phase 6: CuTe Implementation)
    F -->|Ultimate Form| G{Phase 7: CuTe + Swizzling}

Loading

🧠 Memory Hierarchy & Tiling Strategy

To achieve ≥98% of cuBLAS TFLOPS, memory bandwidth must be meticulously managed. We utilize a multi-level tiling strategy:

block-beta
    columns 1
    GlobalMemory["Global Memory (DRAM) - 1000 GB/s"]
    space
    L2Cache["L2 Cache (6-64 MB)"]
    space
    SharedMemory["Shared Memory (SMEM) - 20 TB/s"]
    space
    Registers["Register File - 100 TB/s"]
    space
    TensorCores["Tensor Cores (Math)"]
    
    GlobalMemory --> L2Cache
    L2Cache --> SharedMemory
    SharedMemory --> Registers
    Registers --> TensorCores

Loading

1. Global to Shared Memory: We load blocks of A (size $BM \times BK$) and B (size $BK \times BN$) into Shared Memory to reuse data.
2. Shared Memory to Registers: Warps load fragments from Shared Memory to Registers.
3. Compute: Tensor Cores execute $16 \times 16 \times 16$ mma.sync instructions.

⚔️ Bank Conflict Elimination (Swizzling)

Ampere architectures feature 32 Shared Memory banks (4 bytes wide). Consecutive fp16 elements (2 bytes) can easily trigger bank conflicts during ldmatrix operations.

We use XOR-based address swizzling to mathematically guarantee zero bank conflicts.

Without Swizzling:

• Thread 0 accesses Bank 0
• Thread 16 accesses Bank 0 ❌ (Conflict!)

With XOR Swizzling:

effective_col = col ^ ((row >> 3) & 7) * 8

• Thread 0 accesses Bank 0
• Thread 16 accesses Bank 16 ✅ (Conflict resolved!)

🚀 Setup & Benchmarking (Google Colab)

The project is structured as a PyTorch C++ extension, easily compiled and benchmarked directly in Python.

Quick Start

# 1. Install dependencies
pip install pytest ninja

# 2. Build the extension
pip install -e .

# 3. Verify correctness
pytest tests/test_correctness.py -v

# 4. Run benchmarks against cuBLAS
python benchmarks/bench_vs_cublas.py --backend v1

📊 Performance Tracking

We benchmarked all versions from v1 to v7 against cublasGemmEx on an NVIDIA Tesla T4 GPU (Turing SM75).

📈 TFLOPS vs Matrix Size

This graph showcases the absolute throughput in TeraFLOPS for square matrices of varying sizes:

• v1 (Naive) and v2 (SMEM) are heavily constrained by memory access latency and bank conflicts, peaking around 3 TFLOPS.
• v4 (Async Pipeline) demonstrates a significant leap by introducing asynchronous global-to-shared memory copies, successfully hiding memory latency behind math execution and reaching nearly 9 TFLOPS.
• v5 (SMEM Padding) dominates the benchmark on the Turing architecture by efficiently resolving bank conflicts through stride padding while maintaining native wmma API compatibility, peaking at over 13.3 TFLOPS.
• v7 (CuTe + Swizzle) demonstrates the power of CUTLASS 3.0 abstractions, adapting seamlessly to the SM75 hardware using software UniversalCopy and Turing-specific 16x8x8 MMA atoms to achieve 10.7 TFLOPS.

⚡ Efficiency vs cuBLAS

This graph displays the relative efficiency compared to the highly optimized cuBLAS library:

• By intelligently hiding memory latency and eliminating shared memory bank conflicts, our v5 kernel achieves an impressive 70.2% of cuBLAS efficiency on the T4 GPU.
• The steady progression from 1.9% (v1) to 70.2% (v5) highlights the critical importance of memory hierarchy management and Tensor Core utilization in CUDA kernel optimization.

Version	Description	T4 Efficiency	Bank Conflicts
v1	Naive Implementation	1.9%	N/A
v2	Shared Memory Tiling	15.8%	High
v3	WMMA Tensor Cores	21.0%	High
v4	Async Pipelining	47.6%	High
v5	SMEM Swizzling (Padding)	70.2%	Zero
v6	CuTe Abstractions	45.6%	Zero
v7	CuTe + Swizzle	58.9%	Zero