High-performance custom GPU kernels written in OpenAI's Triton, benchmarked directly against native PyTorch on an NVIDIA H100 SXM (80GB).
This project explores custom fused kernels to bypass PyTorch's eager execution overhead, utilizing cutting-edge Hopper architecture features like FP8 (float8e4m3fn) and fused FlashAttention to maximize memory bandwidth (GB/s) and compute utilization (TFLOPS).
The benchmarks below were executed on a dedicated H100 SXM 80GB environment with locked GPU clocks for hardware consistency.
Memory-bound operations like LayerNorm are restricted by the time it takes to move data between HBM and SRAM. By downcasting the output write from PyTorch's native bfloat16 to the Hopper-specific FP8 format (float8e4m3fn) and calculating a symmetric scaling factor on the fly, we cut the memory I/O footprint by ~17% (the input read and weight/bias reads stay full-width; only the output write shrinks from 2 bytes/elem to 1).
- Peak PyTorch (BF16): 2,422.5 GB/s
- Peak Triton (FP8): 3,531.8 GB/s
- Result: ~45.7% Speedup. The Triton FP8 kernel massively accelerated the layer, showcasing the physical bandwidth advantages of 8-bit precision on the H100.
Standard half-precision LayerNorm benchmark. Native PyTorch executes LayerNorm by calculating the mean, variance, and normalization across multiple separate kernel launches, resulting in redundant trips to HBM. Our custom Triton kernel fuses these operations into a single block-level operation in SRAM.
- Peak PyTorch: 2,696.5 GB/s
- Peak Triton: 3,906.1 GB/s
- Result: ~44.8% Speedup. These are effective bandwidth figures: the numerator counts the bytes the unfused op sequence would move — including per-row weight/bias traffic that is physically served from L2/SRAM — which is why the fused kernel can report more than the H100's ~3.35 TB/s physical HBM peak. Exceeding physical peak is evidence that HBM traffic was eliminated, not that HBM ran faster; physical DRAM traffic (Nsight
dram__bytes) stays below peak.
A custom Triton implementation of FlashAttention. A naive PyTorch attention implementation (manual Q·K^T, softmax, then ·V) scales quadratically with sequence length, causing massive memory bottlenecks when materializing the N x N attention matrix.
- Peak PyTorch Baseline (Seq Len 8192): ~464 TFLOPS
- Result: The fused Triton kernel avoids materializing the attention matrix entirely by relying on online softmax calculations inside the SRAM. As sequence lengths scaled exponentially (512 to 16,384), the Triton kernel bypassed Torch's standard quadratic time and memory limitations entirely, handling ultra-long contexts with negligible overhead.
A fundamentally memory-bound operation, used as a correctness/overhead sanity check rather than an H100 benchmark. This one runs locally on CPU (Triton in interpreter mode vs. eager PyTorch) instead of on the H100, since it's meant to isolate pointer/grid overhead rather than measure hardware bandwidth.
- Result: Triton and Torch matched almost exactly, showing that our base memory pointer alignment and grid instantiation introduce zero meaningful overhead compared to native PyTorch ATen kernels.
├── benchmarks/
│ ├── bench_flash_attn.py # TFLOPS benchmark for Attention
│ ├── bench_layernorm.py # GB/s benchmark for FP16 Fusion
│ ├── bench_layernorm_fp8.py # GB/s benchmark for FP8 Quantization
│ └── bench_relu.py # Baseline memory throughput
├── kernels/
│ ├── flash_attn.py # Block-wise exact attention & online softmax
│ ├── layer_norm.py # Fused parallel reduction normalization
│ ├── layer_norm_fp8.py # Symmetric FP8 scaling and conversion
│ ├── relu.py # Standard activation
│ └── vector_add.py # Elementwise add (Triton 101 warm-up kernel)
├── results/
│ ├── *.csv # Raw performance data output
│ └── *.png # Plotted hardware metrics
└── README.md