In this talk, Tri Dao, the inventor of FlashAttention, explores the landscape of Domain-Specific Languages (DSLs) for GPU kernel optimization. He frames the discussion around the fundamental trade-off between researcher productivity and hardware efficiency, arguing that DSLs are crucial for balancing these needs to advance AI scaling laws. The presentation compares three key DSLs at different abstraction levels: PyTorch (torch.compile) for high-level simplicity, Triton for intermediate control, and the lower-level CuteDSL for maximum hardware access. Using softmax as a primary example, Dao demonstrates how CuteDSL's explicit control over the GPU's memory and thread hierarchy yields superior performance. He presents benchmarks on H100 and Blackwell GPUs, showing that CuteDSL can match or exceed highly optimized libraries like cuBLAS and cuDNN in tasks like matrix multiplication, fused operations (GEMM+SwiGLU), and FlashAttention. The talk concludes with a practical guide on when to use each language, balancing development time against performance gains.
Balancing GPU Performance and Productivity with Domain-Specific Languages
0:01:38 The Core Challenge: To continue the progress of AI scaling laws, it's essential to maximize "intelligence per dollar." This requires balancing algorithmic efficiency (better models) with hardware efficiency (more computations per dollar). DSLs are the key to managing this trade-off.
0:03:35 A Spectrum of DSLs: The talk focuses on three DSLs at different levels of abstraction:
PyTorch (torch.compile): The highest level, extremely easy to use (often just one line of code), and great for researcher productivity.
Triton: An intermediate language that abstracts away complex details like memory management, making GPU programming more accessible than CUDA.
CuteDSL: A lower-level language that exposes the entire GPU hardware hierarchy (threads, warps, blocks, clusters), giving programmers maximum control and performance potential.
0:07:30 Softmax Example: Implementing softmax illustrates the differences:
torch.compile requires a single @torch.compile decorator.
Triton code is concise, requiring about 10-15 lines.
CuteDSL is far more verbose (50+ lines), demanding explicit control over memory vectorization, asynchronous copies, and multi-level reductions.
0:12:53 Performance Gains from Control: Benchmarks on an H100 GPU show that the extra effort with CuteDSL pays off. It achieves up to 15-30% better performance on softmax for both very small and very large inputs by leveraging warp-level and cluster-level reductions, features not easily accessible in higher-level languages.
0:14:44 Matching State-of-the-Art Libraries: In matrix multiplication (GEMM) benchmarks, a CuteDSL implementation matches and sometimes even outperforms NVIDIA's highly optimized cuBLAS library on H100 and Blackwell GPUs. This proves that Python-embedded DSLs can achieve peak hardware performance.
0:17:53 The Power of Kernel Fusion: For common operations like GEMM followed by a SwiGLU activation, fusing them into a single CuteDSL kernel provides a 7-15% speedup compared to the standard approach of using two separate kernels (e.g., cuBLAS + Triton), which is what torch.compile would do.
0:19:11 Pushing the Frontier with FlashAttention: A CuteDSL implementation of FlashAttention on the new Blackwell GPU is 15-20% faster than NVIDIA's closed-source cuDNN library, demonstrating the power of open and customizable kernels for cutting-edge algorithms.
Onboarding Time:torch.compile takes hours to days to learn, Triton takes days to weeks, and CuteDSL takes weeks to months due to the required understanding of GPU internals.
When to Use What:
Start with torch.compile: It's the easiest and is great for simple fusions.
Move to Triton: Use for more complex memory-bound kernels where more control is needed.
Use CuteDSL for Peak Performance: Reserve it for critical compute-bound kernels (like GEMM or attention) where custom fusion and squeezing out every bit of performance is necessary.