1,001 Ways to Accelerate Python with CUDA Kernels | NVIDIA GTC 2025

1,001 Ways to Accelerate Python with CUDA Kernels

Speakers: Leo Fang, Python CUDA Tech Lead, NVIDIA

Conference: NVIDIA GTC 2025

YouTube: https://www.youtube.com/watch?v=_XW6Yu6VBQE

Overview

In this comprehensive GTC 2025 presentation, Leo Fang, NVIDIA's Python CUDA Tech Lead, unveils a spectrum of innovative approaches for accelerating Python applications by leveraging CUDA kernels. The talk, part of a series dedicated to CUDA Python, specifically focuses on the art and science of authoring high-performance kernels from within the Python ecosystem. Fang emphasizes NVIDIA's overarching mission to foster a robust and interoperable CUDA Python environment, where developers can seamlessly mix and match various Python packages and CUDA functionalities. Using the pedagogical example of segmented reduction—a common operation in data processing and machine learning—Fang meticulously walks through the evolution from traditional C++ CUDA kernel development to a rich landscape of Python-centric tools and programming models, all while striving to maintain C++-level performance. This talk is crucial for Python developers aiming to push the boundaries of performance on NVIDIA GPUs without having to abandon their preferred language or grapple with the complexities of C++ compilation workflows.

The core challenge addressed is the historical friction between Python's productivity and CUDA's performance, traditionally requiring developers to dip into C++ for custom, high-speed GPU code. Fang demonstrates how modern compiler technologies and new NVIDIA libraries are dissolving this barrier, enabling Python developers to author, compile, and deploy CUDA kernels with unprecedented ease and efficiency. From JIT compilation of C++ kernels from Python strings to entirely new cooperative programming models and direct access to highly optimized device libraries, the presentation showcases a future where Python is a first-class citizen in the realm of high-performance GPU computing. The ultimate goal is to empower Python users to achieve near "speed of light" performance for their custom algorithms, iterating rapidly within their familiar environment.

Background

▶ Watch: Introduction to CUDA Python kernel development (0:00)

The proliferation of AI and machine learning has dramatically increased the demand for high-performance computing, with GPUs becoming indispensable accelerators. Python, with its rich ecosystem of data science and ML libraries (like NumPy, PyTorch, TensorFlow, JAX, CuPy), has emerged as the language of choice for many researchers and practitioners. However, a perennial challenge has been bridging the gap between Python's ease of use and the raw performance offered by NVIDIA's CUDA platform. Traditionally, achieving maximum performance for custom GPU operations required writing kernels in C++ using the CUDA programming model. This approach introduced significant friction:

1. Complexity of CUDA C++: Developers needed deep knowledge of GPU architecture, memory hierarchies (global, shared, register), thread and block management, and synchronization primitives.
2. C++ Compilation Workflow: Building C++ CUDA kernels typically involves host compilers (GCC, MSVC), complex build systems (CMake), and meticulous management of template specializations and library dependencies. This often leads to large binary sizes and slow compile times, especially when supporting multiple data types.
3. Deployment Challenges: Shipping C++-based CUDA libraries to Python users meant requiring users to install specific host compilers and manage intricate build environments, which is often not "pip installable" or "conda installable" friendly.
4. Interoperability: Integrating custom C++ CUDA kernels into Python frameworks often involved cumbersome wrapper code, adding overhead and reducing development velocity.

This problem is particularly acute for tasks like segmented reduction, which serves as the pedagogical example in the talk. Segmented reduction involves performing a reduction (e.g., sum, min, max) independently over non-overlapping, contiguous segments within a larger 1D data buffer. While seemingly simple, an efficient GPU implementation requires careful management of parallel threads, shared memory, and reduction strategies (like tree reduction) to avoid race conditions and maximize throughput. Prior to recent advancements, Python users often resorted to existing optimized libraries (if available) or faced the daunting task of writing and integrating a custom C++ CUDA kernel, undermining Python's promise of productivity. The need for a more Python-native, performant, and developer-friendly way to author and deploy CUDA kernels became increasingly evident to unlock the full potential of GPU acceleration for the Python community.

Key Findings

▶ Watch: Traditional C++ CUDA kernel implementation overview (2:30)

The talk reveals several pivotal advancements and programming paradigms that collectively redefine how Python developers can leverage CUDA kernels, effectively bridging the historical gap between Python's productivity and C++'s performance on GPUs. The overarching key finding is that Python can now achieve C++-level CUDA kernel performance with significantly enhanced developer productivity and a streamlined workflow, largely thanks to modern JIT compilation techniques and the introduction of Python-native device libraries.

Specific key findings include:

1. Elimination of C++ Host Compiler Dependency: Modern CUDA Python tools, powered by NVRTC, NVGLink, and NVVM, enable Just-In-Time (JIT) compilation of CUDA kernels. This crucial development means Python projects no longer need to depend on complex C++ host compilers (like GCC or MSVC) for building GPU code, simplifying deployment and improving user experience.
2. Runtime Template Specialization in Python: What were traditionally C++ compile-time template parameters (e.g., data types, block sizes) can now be specified and instantiated at Python runtime. This allows for dynamic kernel generation tailored to specific input arguments, leading to smaller binary sizes and faster compilation times compared to pre-instantiating all possible C++ template variations.
3. Viable Python Device Libraries: The introduction of NVGLink (starting with CUDA 12.0) and the use of LTO IR (LLVM Bitcode) finally enables the creation and distribution of high-performance Python-defined device libraries. This allows for the abstraction of complex, optimized GPU primitives (like block-wide reductions or loads) into reusable Python functions, significantly simplifying user kernel code and promoting modularity without performance overhead.
4. New Cooperative Programming Models: NVIDIA is introducing novel programming models, notably CUDA Tile (CouTile), which shifts the abstraction from individual threads to "tile blocks." This NumPy-like approach simplifies kernel authoring by handling intricate thread management and pointer arithmetic, allowing developers to focus on data operations on chunks of memory with built-in cooperative primitives.
5. Direct Python Access to Optimized GPU Primitives: Projects like CUDA Cooperative (based on CUB), NVMath Python, and Cutlass Python provide Pythonic interfaces to highly optimized C++ device libraries and low-level Tensor Core controls. This grants Python developers unprecedented access to "speed of light" performance-tuned algorithms and hardware features without writing a single line of C++.
6. Emphasis on Avoiding Kernel Writing: A practical, albeit counter-intuitive, key finding is the recommendation to avoid writing custom kernels whenever possible. The talk highlights that a growing ecosystem of high-level Python libraries (e.g., CUDA Parallel, NVMath host APIs, Array API standard-compliant libraries like CuPy, JAX, PyTorch) can often solve common problems with optimal performance, offering greater portability and less development burden than custom kernel development.

These findings collectively represent a significant leap forward in the CUDA Python ecosystem, empowering a broader range of developers to harness the full power of NVIDIA GPUs with Pythonic elegance and efficiency.

Technical Deep Dive

▶ Watch: Detailed C++ CUDA kernel logic and shared memory (4:00)

The talk meticulously dissects various technical approaches to writing and deploying CUDA kernels from Python, moving from low-level C++ integration to high-level, Python-native programming models. The pedagogical example used throughout is segmented reduction, where an input 1D buffer containing multiple segments (e.g., red, blue, yellow) needs to have a sum (or other reduction) computed for each segment independently.

1. The Traditional C++ CUDA Kernel Approach

Leo Fang begins by illustrating the complexity of a hand-written C++ CUDA kernel for segmented reduction. This kernel assigns one segment to a CUDA thread block. Key components include:

• Template Parameters: The kernel is templated over data types, offset types, block size, and items per thread, allowing compile-time optimization.
• Shared Memory Allocation: Essential for inter-thread communication within a block, particularly for parallel reduction.
• Index Computation: Each thread computes its unique index within the block and its segment index within the batch.
• Partial Aggregation: Threads load data from global memory into registers, performing initial partial sums.
• Shared Memory Reduction: After loading and partial aggregation, threads synchronize (__syncthreads()) and then perform a tree reduction using shared memory. This involves a pattern where half the active threads write their results, then a quarter, and so on, until a single thread holds the final segment sum.
• Global Memory Write: The final result from the reducing thread is written back to global memory.

This approach, while performant, demands deep understanding of the CUDA programming model, GPU memory hierarchy, and C++ compilation workflows. It requires explicit management of thread blocks, shared memory, and synchronization, making it labor-intensive and error-prone.

2. JIT Compilation for Python Productivity

To overcome C++'s compilation hurdles, Fang introduces the concept of Just-In-Time (JIT) compilation within Python.

• Problem: C++ kernels require host compilers (GCC, MSVC) and pre-instantiation of templates, leading to large binaries and slow build times.
• Solution: NVIDIA's online compilation toolchains (NVRTC, NVGLink, NVVM) eliminate the need for host compilers. They compile user code directly to GPU machine code or intermediate representations at runtime.
• CUDA Core: This new NVIDIA project provides Python bindings to CUDA APIs, including memory management, streams, events, and crucially, the online compilers. Developers can load C++ CUDA kernel code as a Python string, specify template specializations (e.g., float for data, int for offsets, specific block size) at runtime, compile the specialized kernel, and launch it. This means C++ compile-time information (templates, constexprs) becomes Python JIT-time information, enabling dynamic kernel generation and faster iteration. CUDA Core supports CUDA 11 and 12, is pip/conda installable, and works on Linux and Windows.

3. Python-Native Kernel Authoring with Numba CUDA

While CUDA Core allows C++ kernels to be JIT-compiled, many Python developers prefer to write kernels entirely in Python.

• Numba CUDA: This open-source project (developed by Anaconda, now being decoupled from Numba's release cycle) allows developers to write CUDA kernels using a Pythonic syntax. The example of squaring array elements is shown, demonstrating a cleaner syntax (e.g., **2 for power) compared to C++.
• Segmented Reduction in Numba CUDA: The complex segmented reduction kernel can be rewritten in Numba CUDA. While the syntax is Pythonic, the underlying logic (shared memory management, parallel tree reduction, thread indexing) remains largely the same as the C++ version. This highlights that Numba CUDA simplifies syntax but still requires deep CUDA programming knowledge for optimal performance.

4. Enabling Python Device Libraries with NVGLink

A significant limitation for Python CUDA kernel development was the difficulty of creating reusable, high-performance device functions (like BlockReduce).

• Historical Problem: Python device libraries were either C++-based (overhead) or relied on PTX (lacked minor version compatibility, making them fragile).
• NVGLink (CUDA 12.0+): This library solves the problem by linking multiple pieces of LTO IR (LLVM Bitcode). Python device functions can now be compiled into LTO IR, linked with user kernels, and inlined, producing performance equivalent to C++ counterparts with no function call overhead. This "extends the life of a Numba kernel" by allowing it to compose with optimized device libraries.

5. High-Performance Primitives via CUDA Cooperative

Leveraging NVGLink, NVIDIA is developing CUDA Cooperative, a project exposing high-performance C++ parallel algorithms (specifically from the CUB library) to Numba users.

• Concept: Developers define device functions for operations like loading and reduction in Python. CUDA Cooperative then generates the necessary LTO IR under the hood.
• Example (Segmented Reduction): The previously complex Numba CUDA kernel is drastically simplified. Instead of hand-writing shared memory loading and tree reduction, the user kernel now calls cub.BlockLoad and cub.BlockReduce (or similar primitives) as device functions.
• Benefits: Dramatically cleaner, more manageable user kernels. Supports arbitrary Python-defined data types and operators. Results in a small, open-source package. This library is still under development but is publicly available for cloning and building.

6. Leveraging Tensor Cores with NVMath Python

Segmented reduction can be re-framed as a matrix-vector multiplication for specific sparse matrix structures.

• Approach: Think of the segments as rows of a 2D matrix, and multiply by a vector of ones.
• NVMath Python: This library provides Pythonic access to the NVMath device library. Developers declare device functions using nvmath.device.mat which helps allocate optimal shared memory. Inside the kernel, shared memory is allocated for input (A), the vector of ones (B), and output (C). The nvmath.device.mat API is then called to compute C = A * B.
• Tensor Core Utilization: A key advantage is that nvmath.device.mat automatically leverages Tensor Cores for accelerated computation, delivering significant performance gains.

7. The New CUDA Tile Programming Model (CouTile)

Recognizing the pervasive "cooperative" nature of GPU algorithms, NVIDIA is introducing a new programming model.

• Concept: CUDA Tile (CouTile) shifts the programming paradigm from individual threads to tile blocks. Developers operate on "tiles" (chunks of data loaded into registers/shared memory) using NumPy-like operations.
• Benefits: Decouples thread management and pointer arithmetic from data loading and computation. Simplifies common patterns like sort, scan, reduction, and matrix multiplication. The compiler handles the intricate details of thread coordination.
• Example: For squaring array elements, a CouTile kernel might simply load a tile, apply a square operation, and store the result, abstracting away individual thread logic.
• NVIDIA Warp Integration: NVIDIA Warp, a differential programming framework for simulations and AI, has also integrated this tile programming model. Warp additionally offers automatic synthesis of backward kernels for gradient computation, now open-source under Apache 2.0.

8. Full Tensor Core Control with Cutlass Python

For "ninjas" who demand ultimate control over GPU hardware:

• Cutlass Python: This upcoming project will provide Python support for Cutlass, NVIDIA's highly optimized C++ template library for linear algebra primitives, especially for Tensor Core operations.
• Benefits: Exposes all the low-level knobs for fine-grained control over Tensor Cores from Python. Significantly reduces compile times compared to C++ Cutlass, allowing rapid iteration while achieving "speed of light" performance.

9. Community and Existing Tools

Fang also acknowledges the broader ecosystem:

• CuPy: Offers high-performance array operations and custom kernel capabilities (cupy.RawKernel, cupy.ReductionKernel).
• CuJIT: A JIT compiler for CUDA kernels in Python.
• Triton: OpenAI's DSL for writing highly optimized GPU kernels, featuring a block-based programming model.

The talk concludes this section by emphasizing the vast array of choices, both first-party and third-party, reinforcing the notion that the CUDA Python ecosystem is thriving and diverse.

Experimental Setup & Results

▶ Watch: Staying productive: Python's advantage for CUDA kernel development (6:58)

This talk primarily serves as an architectural overview and a guide to available tools and programming models within the CUDA Python ecosystem, rather than presenting novel experimental results or benchmarks for a specific, new algorithm. The core objective is to demonstrate how developers can achieve high performance and productivity, not to provide comparative performance numbers for new findings.

Therefore, the presentation does not detail a specific experimental setup with datasets, baselines, or quantitative headline numbers. Instead, it uses segmented reduction as a pedagogical example to illustrate the process of kernel development across different tools and paradigms.

• Pedagogical Example: The segmented reduction problem is defined clearly, showing input segments (red, blue, yellow) and their desired summed outputs (10, 26, etc.). This example is consistent across C++, Numba CUDA, CUDA Cooperative, and NVMath Python demonstrations.
• Performance Claims: While specific benchmark numbers are not provided, the talk consistently asserts that the goal of these new tools and programming models is to achieve C++-level performance from Python. For instance, NVGLink's ability to link LTO IR and inline device functions is highlighted as enabling kernel generation "as performant as the C++ counterpart." Similarly, Cutlass Python is stated to allow developers to "write a speed of light kernel to get maximum performance out of our GPUs."
• Hardware: The underlying hardware is implicitly NVIDIA GPUs, given the context of CUDA and GTC. Specific GPU types (e.g., A100, H100) are not mentioned in the context of results, but the mention of Tensor Cores (via NVMath Python and Cutlass Python) implies modern NVIDIA GPU architectures.
• Referral to Other Talks: For detailed performance analysis and benchmarks, Fang explicitly refers attendees to other GTC talks by his colleagues (e.g., Georgie and Bryce for C++ performance, Stephen Jones for CUDA Tile new features).

In essence, the "results" presented in this talk are qualitative: the demonstration of cleaner code, faster iteration, and the potential for C++-equivalent performance through Python-native development, rather than specific quantitative performance metrics. The focus is on the enabling technologies and programming models that make such performance achievable.

Practical Implications

▶ Watch: Online compilation machinery for Python CUDA kernels (7:39)

The advancements detailed in this talk have profound practical implications for a wide array of professionals in the AI/ML and high-performance computing space:

1. For Python Practitioners (Data Scientists, ML Engineers):

• Increased Productivity and Faster Iteration: The ability to write, compile, and debug CUDA kernels entirely within Python, often with Pythonic syntax (Numba CUDA, CUDA Tile), dramatically speeds up the development cycle. This means less context switching between languages and build systems.
• Access to Cutting-Edge Performance: Tools like CUDA Cooperative, NVMath Python, and Cutlass Python democratize access to highly optimized, low-level GPU primitives and Tensor Core functionality. Practitioners can achieve near "speed of light" performance for custom operations without needing to be C++ experts.
• Simplified Deployment: JIT compilation removes the dependency on C++ host compilers, making Python projects with custom CUDA kernels easier to package, distribute, and install via standard tools like pip and conda.
• "Avoid Kernel Writing" Mindset: The emphasis on utilizing existing high-level libraries (CUDA Parallel, NVMath host APIs, Array API standard libraries like CuPy, JAX, PyTorch) encourages practitioners to first seek out optimized, portable solutions, reducing the burden of custom kernel development. This promotes more robust and maintainable codebases.

1. For Infrastructure Teams and MLOps Engineers:

• Streamlined CI/CD: The removal of host compiler dependencies simplifies continuous integration and deployment pipelines for Python projects that include custom GPU code. Build environments become lighter and more consistent.
• Reduced Support Matrix Complexity: While the talk acknowledges the need to consider support matrices (Python, CUDA, driver versions, OS, CPU/GPU architectures), the JIT compilation approach can reduce the need to pre-compile for every possible permutation, potentially simplifying maintenance.
• Smaller Binaries: Dynamic kernel generation based on runtime needs (e.g., specific data types) can lead to smaller deployed binaries compared to C++ libraries that pre-instantiate all possible template combinations.

1. For Model Builders and Researchers:

• Custom Op Flexibility: Researchers developing novel algorithms or custom layers can implement them efficiently on GPUs directly in Python, allowing for rapid prototyping and integration into existing ML frameworks.
• Differentiable Kernels: Frameworks like NVIDIA Warp, with its tile programming model and automatic backward kernel synthesis, are invaluable for researchers working on differentiable simulations, physics-informed neural networks, or custom gradient computations.

Tradeoffs and Limitations:

• Learning Curve: While Pythonic syntax is easier, mastering optimal CUDA kernel performance still requires understanding GPU architectures and parallel programming principles, even with higher-level abstractions like CUDA Tile or CUDA Cooperative.
• Tool Choice Complexity: With "1,001 ways," choosing the "best" tool remains a non-trivial decision, dependent on factors like existing dependencies, performance targets, support matrix, packaging requirements, and organizational preferences. The talk explicitly states "it depends."
• Early Stage Libraries: Some of the highly promising libraries, like CUDA Cooperative and CUDA Tile, are still under active development or in early release stages. Practitioners might need to build from source or anticipate API changes.
• Debugging: While iteration is faster, debugging complex GPU kernels, especially those leveraging low-level features or new programming models, can still be challenging.

In summary, these advancements signify a maturation of the CUDA Python ecosystem, moving beyond simple bindings to offering a comprehensive suite of tools and programming models that enable Python developers to achieve high-performance GPU acceleration with unprecedented ease and efficiency. The strategic advice to leverage existing solutions whenever possible further empowers the community to build robust and performant AI/ML applications.

Key Takeaways

• Python is a First-Class Language for CUDA Kernel Development: Modern compiler technologies (NVRTC, NVGLink, NVVM) and new NVIDIA libraries now enable Python developers to write, compile, and deploy high-performance CUDA kernels without relying on C++ host compilers, achieving C++-level performance.
• JIT Compilation Revolutionizes Workflow: Just-In-Time (JIT) compilation from Python allows dynamic kernel specialization at runtime, leading to faster iteration cycles, smaller binary sizes, and simpler deployment for CUDA-accelerated Python projects.
• Python Device Libraries are Now Viable: With CUDA 12.0 and NVGLink leveraging LTO IR, Python-defined device functions can be compiled and linked into user kernels with zero overhead, enabling the creation of reusable, high-performance Python device libraries like CUDA Cooperative.
• New Programming Models Simplify GPU Development: CUDA Tile (CouTile) introduces a higher-level, NumPy-like programming model that operates on "tile blocks," abstracting away complex individual thread management and pointer arithmetic, making kernel authoring more intuitive and productive.
• Direct Access to GPU Hardware and Optimized Libraries: Projects like NVMath Python and the upcoming Cutlass Python provide Pythonic interfaces to hardware features (e.g., Tensor Cores) and highly optimized C++ libraries, allowing Python developers to extract maximum GPU performance.
• Prioritize Existing Solutions: The most practical advice is to avoid writing custom kernels if an optimized solution already exists in the vast and growing ecosystem of high-level Python libraries (e.g., CUDA Parallel, CuPy, NVMath host APIs, Array API standard-compliant libraries).

About the Speaker(s)

Leo Fang is the Python CUDA Tech Lead at NVIDIA. In this role, he is at the forefront of developing tools and frameworks that bridge the gap between Python's productivity and CUDA's high-performance capabilities. His work focuses on enhancing the entire CUDA Python ecosystem, enabling developers to seamlessly integrate GPU acceleration into their Python workflows. Fang is deeply involved in projects that leverage modern compiler technologies and introduce innovative programming models to simplify CUDA kernel authoring for the Python community.