1. GEMM (General Matrix Multiplication):
- This is the basic operation: C = A × B (matrix multiplication)
- Fundamental operation in deep learning, especially transformers
- Core computation in attention mechanisms, linear layers, etc.
2. Triton:
- A programming language for writing GPU kernels
- Lets you write your own custom GEMM implementation
- You control memory layout, tiling, etc.
- Example use: When you need a very specific matrix operation
3. hipBLASLt:
- A specialized library just for matrix operations
- Pre-built, highly optimized GEMM implementations
- Focuses on performance for common matrix sizes
- Example use: When you need fast, standard matrix multiplication
4. Transformer Engine:
- NVIDIA's specialized library for transformer models
- Automatically handles precision switching (FP8/FP16/FP32)
- Optimizes GEMM operations specifically for transformer architectures
- Includes specialized kernels for attention and linear layers
- Example use: When building large language models
The relationship:
```
Transformer Model
↓
Transformer Engine
↓
GEMM Operations (can be implemented via:)
↓
hipBLASLt / Triton / Other libraries
↓
GPU Hardware
```
the same matrix multiplication would be implemented using different approaches:
1. Basic GEMM Operation (what we want to compute):
```python
# C = A × B
# Where A is (M×K) and B is (K×N)
```
2. Using Triton (Custom implementation):
```python
@triton.jit
def matmul_kernel(
a_ptr, b_ptr, c_ptr, # Pointers to matrices
M, N, K, # Matrix dimensions
stride_am, stride_ak, # Memory strides for A
stride_bk, stride_bn, # Memory strides for B
stride_cm, stride_cn, # Memory strides for C
BLOCK_SIZE: tl.constexpr,
):
# Get program ID
pid = tl.program_id(0)
# Calculate block indices
block_i = pid // (N // BLOCK_SIZE)
block_j = pid % (N // BLOCK_SIZE)
# Load blocks from A and B
a = tl.load(a_ptr + ...) # Load block from A
b = tl.load(b_ptr + ...) # Load block from B
# Compute block multiplication
c = tl.dot(a, b) # Matrix multiply
# Store result
tl.store(c_ptr + ..., c)
```
3. Using hipBLASLt:
```cpp
// Initialize hipBLASLt
hipblasLtHandle_t handle;
hipblasLtCreate(&handle);
// Define matrix layout
hipblasLtMatrixLayout_t matA, matB, matC;
hipblasLtMatrixLayoutCreate(&matA, HIPBLAS_LT_R_16F, M, K, M);
hipblasLtMatrixLayoutCreate(&matB, HIPBLAS_LT_R_16F, K, N, K);
hipblasLtMatrixLayoutCreate(&matC, HIPBLAS_LT_R_16F, M, N, M);
// Execute GEMM
hipblasLtMatmul(
handle,
matmulDesc,
&alpha, // Scale factor
A, matA, // Input matrix A
B, matB, // Input matrix B
&beta, // Scale factor
C, matC, // Output matrix C
workspace, // Temporary workspace
streams // CUDA stream
);
```
4. Using Transformer Engine:
```python
import transformer_engine.pytorch as te
# Create TE layers
linear = te.Linear(in_features, out_features)
# Automatic precision handling
with te.fp8_autocast():
output = linear(input) # Internally uses optimized GEMM
```
Key differences:
1. Triton: You control everything (memory, blocks, compute)
2. hipBLASLt: Pre-optimized, you just call it
3. Transformer Engine: High-level, handles precision automatically
Performance comparison (general case):
```
Speed: hipBLASLt > Transformer Engine > Custom Triton
Flexibility: Triton > hipBLASLt > Transformer Engine
Ease of use: Transformer Engine > hipBLASLt > Triton
```
