NVIDIA’s Blackwell Architecture with Triton

Modern AI workloads rely heavily on matrix multiplication and attention mechanisms, which form the computational backbone of machine learning models. While NVIDIA cuDNN offers highly optimized implementations and frameworks like CUTLASS, which provide deep customization, many developers seek a balance between performance and flexibility. The open-source Triton compiler, now enhanced for the NVIDIA Blackwell architecture, addresses this need by exposing Blackwell’s advanced features through an intuitive programming model.

This collaboration between OpenAI and NVIDIA enables developers to leverage Blackwell’s capabilities with Triton’s Python-based compiler, ensuring easy access to high-performance AI computing.

Performance Enhancements on NVIDIA Blackwell

The NVIDIA Blackwell architecture brings substantial improvements in raw computing power, focusing on two key areas:

• Optimized Matrix Multiplication, including support for new precision formats.
• Flash Attention Acceleration, delivering significant speedups for transformer models.

Matrix Multiplication: Leveraging New Tensor Cores

Blackwell introduces a new Tensor Core designed for improved throughput and energy efficiency. Developers can automatically exploit these enhancements by extending Triton’s Matrix Multiply-Accumulate (MMA) pipelining. This required optimizing memory access patterns and compiler transformations to ensure efficient compute/data-movement overlap.

As a result, Triton achieves near-optimal performance for FP8 and FP16 General Matrix Multiplication (GEMM) operations, applying optimizations automatically to kernels using Triton’s tl.dot primitive. Benchmarks show significant speedups on Blackwell GPUs, outperforming previous generations such as NVIDIA Hopper.

Flash attention, a fundamental operation in modern transformer architectures, sees up to 1.5x performance improvement on FP16 attention workloads compared to NVIDIA Hopper. These optimizations allow developers to transition to Blackwell with zero code changes, as Triton seamlessly integrates these improvements into existing implementations.

Introducing New Precision Formats

NVIDIA Blackwell introduces block-scaled floating point formats, including OCP’s microscaling formats, which Triton now unlocks for hardware acceleration. These formats provide:

• Higher average precision compared to non-native block-scaling techniques.
• Accelerated GEMM operations using MXFP8 and MXFP4, enhancing both performance and precision.

Notably, MXFP4 doubles the performance of FP8 and MXFP8 GEMMs, offering a compelling precision-performance trade-off. Developers can explore these advancements through Triton’s new block-scaled floating point support tutorials.

Future Improvements and Community Engagement

While Triton now fully supports NVIDIA Blackwell, ongoing improvements will enhance usability:

• Optimizing Sub-Byte Formats: Handling MXFP4 and similar data formats more efficiently.
• Enhancing GEMM Utilization: Improving performance for smaller GEMM_K values through automatic warp-specialization in the compiler.

NVIDIA and OpenAI continue to refine these features, with more details to be shared at NVIDIA GTC 2025 on March 17.