@PyTorch: In this post, you’ll learn how to use the PyTorch Torch Inductor compiler and kernel fusion to improve memory bandwidth…
摘要
NVIDIA and PyTorch explain how kernel fusion in CUDA improves GPU memory bandwidth by combining multiple operations into a single kernel, reducing round-trips through global memory and kernel launch overhead.
查看缓存全文
缓存时间: 2026/08/04 20:15
In this post, you’ll learn how to use the PyTorch Torch Inductor compiler and kernel fusion to improve memory bandwidth and reduce kernel launch overhead.
A common bottleneck when writing GPU code is that GPU compute is so fast that even high-bandwidth device memory doesn’t use the GPU kernel fully. Kernel fusion addresses this by combining multiple GPU operations into a single device kernel, so intermediate results don’t need to round-trip through global memory or require separate kernel launches.
We can describe a computation in plain Python and let the compiler generate the kernel. That is what torch.compile does. You give it a function, it traces what the tensors are doing, and Torch Inductor, PyTorch’s compiler, generates a kernel that runs on the GPU.
Read the full post: https://bit.ly/45Emn5g
Kernel Fusion in NVIDIA CUDA: Optimizing Memory Traffic and Launch Overhead
Source: https://developer.nvidia.com/blog/kernel-fusion-in-nvidia-cuda-optimizing-memory-traffic-and-launch-overhead/ There are many ways to optimize code for GPUs. In this post, you’ll learn how kernel fusion can improve memory bandwidth and reduce kernel launch overhead, along with multiple ways to apply it in NVIDIA CUDA code.
A common bottleneck when writing GPU code is that GPU compute is so fast that even high-bandwidth device memory doesn’t use the GPU kernel fully. Kernel fusion addresses this by combining multiple GPU operations into a single device kernel, so intermediate results don’t need to round-trip through global memory or require separate kernel launches.
This post is about fusing kernel bodies so intermediate results stay in registers, and memory transfers are reduced. CUDA Graphs provide another kind of fusion, but at a different layer. A graph captures a sequence of kernel launches, memory copies, and synchronizations into one reusable object that the host can dispatch with a single call. It doesn’t fuse kernel bodies.
Kernel fusion in practicehttps://developer.nvidia.com/blog/kernel-fusion-in-nvidia-cuda-optimizing-memory-traffic-and-launch-overhead/#kernel_fusion_in_practice
Kernels inside a graph run separately and pass intermediate results through global memory. Wrapping a naive baseline in a graph shaves microseconds off the host side, but the one GiB round-trip through the intermediate buffer remains unchanged. The two approaches are complementary and can be used together.
Let’s take a simple example:sum\(abs\(x\)\). We read an array, take the absolute value of each element, and return the sum.
A naive implementation uses two kernels: one to computeabsinto a temporary intermediate buffer of the same size as the input, and another to reduce that buffer to a single number. This works, but it’s slow in a way that shows the broader problem.
template <typename Config>
__global__ void abs_kernel(Config config,
cuda::std::span<const float> in,
cuda::std::span<float> tmp)
{
/* Base SIMT implementation using fsabs */
}
template <typename Config>
__global__ void sum_kernel(Config config,
cuda::std::span<const float> tmp,
cuda::std::span<float> out)
{
/* CUB block-reduce + atomic add into out[0] */
}
int main()
{
/* device, stream, and buffers via cuda::make_buffer
* See the CCCL Runtime post for the full setup. */
auto config = cuda::distribute<BLOCK_THREADS>(in.size());
cuda::launch(stream, config, abs_kernel<decltype(config)>, in, tmp);
cuda::launch(stream, config, sum_kernel<decltype(config)>, tmp, out);
}
The C++ examples featured throughout this post use the latest NVIDIA CCCL runtime interfaces, includingcuda::std::span,cuda::launch, andcuda::make\_buffer—which debuted in CUDA 13.2. To explore these primitives further, refer toCCCL Runtime: A Modern C++ Runtime for CUDA.
If we look at an NVIDIANsight Systemsprofile of the full implementation, it looks like this:
Figure 1. Nsight Systems profile of non-fused kernels for sum(abs(x)). Each kernel is executed independently, and an intermediate result is needed
Manual kernel fusionhttps://developer.nvidia.com/blog/kernel-fusion-in-nvidia-cuda-optimizing-memory-traffic-and-launch-overhead/#manual_kernel_fusion
The first thing we could do is rewrite these two kernels as a singlesum\_abs\_kernel. This removes the intermediate buffer, which exists only to pass data between kernels. If both operations live in the same kernel, the buffer disappears, and the performance of the operation improves.
template <typename Config>
__global__ void sum_abs_kernel(Config config,
cuda::std::span<const float> x,
cuda::std::span<float> out)
{
using BlockReduce = cub::BlockReduce<float, BLOCK_THREADS>;
__shared__ typename BlockReduce::TempStorage temp_storage;
// Grid-stride loop. abs() is computed inline -- no tmp buffer.
float thread_sum = 0.0f;
const auto tid = cuda::gpu_thread.rank(cuda::grid, config);
const auto stride = cuda::gpu_thread.count(cuda::grid, config);
for (size_t i = tid; i < x.size(); i += stride) {
thread_sum += fabsf(x[i]);
}
// Block-level reduction in shared memory, via CUB.
float block_sum = BlockReduce(temp_storage).Sum(thread_sum);
// One thread per block atomically accumulates into the global result.
if (threadIdx.x == 0) {
cuda::atomic_ref<float, cuda::thread_scope_device> r(out[0]);
r.fetch_add(block_sum, cuda::memory_order_relaxed);
}
}
int main()
{
/* device, stream, and buffers -- see the CCCL Runtime post. */
// Fixed grid of 1024 blocks for the grid-stride pattern.
auto config = cuda::make_config(cuda::make_hierarchy(
cuda::grid_dims(1024),
cuda::block_dims<BLOCK_THREADS>()));
cuda::launch(stream, config, sum_abs_kernel<decltype(config)>, x, out);
}
fabsf\(x\[i\]\)is computed inline, in a register, instead of being read from an intermediate buffer that another kernel wrote.- A grid-stride loop lets each thread accumulate multiple elements into a single register (
thread\_sum), so partial sums never touch global memory. cub::BlockReducehandles the per-block reduction in shared memory, andcuda::atomic\_reflets one thread per block add its partial sum to the global result.
If we look at an Nsight Systems profile of the fused implementation, it collapses accordingly:
Figure 2. Nsight Systems profile of a manually fused kernel. Asum\_abs\_kernelis a single kernel that does the whole operation
MetricNaiveManual fusionKernels Launched21Intermediate bufferYesNoBytes moved through global memory3 GB1 GBTime3.51 ms = 2.28ms + 1.23 ms1.18 msEffective memory bandwidth855 GiB/s851 GiB/sSpeedup vs. naive–3 timesTable 1. Benchmark for the manually fused kernelThese numbers are for an NVIDIA GeForce RTX 4090. The memory bandwidth sits at ~850 GiB/s—roughly 90% of its theoretical 1,008 GB/s peak.
We’re bandwidth-bound as the GPU is doing as much memory work as physics enables. The fused version is faster because there’s less memory work to do.
A hand-written CUDA C++ kernel like this is sometimes the right solution, because you have full control and get as close as possible to peak performance. The drawback is the effort required to write the kernel, along with the high maintenance cost and limited portability across GPU architectures.
Device-side libraries such as CUB,libcu\+\+,nvmath\.device, and NVSHMEM can simplify manual fusion by providing reusable primitives for FFTs, matrix multiplications, reductions, scans, sorts, communication, and more.
Implicit kernel fusionhttps://developer.nvidia.com/blog/kernel-fusion-in-nvidia-cuda-optimizing-memory-traffic-and-launch-overhead/#implicit_kernel_fusion
Writing your own kernel works. In this case, the operation is essentially “applyingabswhile reducing.” Instead of expressing that manually in CUDA C++, we can describe the computation in plain Python and let the compiler generate the kernel.
That is whattorch\.compiledoes. You give it a function, it traces what the tensors are doing, and Torch Inductor, PyTorch’s compiler, generates a kernel that runs on the GPU.
import torch
x = torch.randn(N, device="cuda")
def sum_abs(t):
return t.abs().sum()
compiled = torch.compile(sum_abs)
result = compiled(x)
By wrapping thesum\_absfunction withtorch\.compilewe get a fused kernel. In the eager version (withouttorch\.compile) we get what we expect: two kernels, one element-wise (abs) and one reduce (sum).
Figure 3. PyTorch profile viewed in Perfetto.dev screenshot for a single eager PyTorch operation. Two kernels are shown, one per operation
In the compiled version, we also get two kernels:triton\_red\_fused\_abs\_sum\_0andtriton\_red\_fused\_abs\_sum\_1.
Why? Because that is what the Inductor compiler generates.
Figure 4. PyTorch profile viewed in Perfetto.dev screenshot for the compiled PyTorch. These two kernels are generated by PyTorch, even if the operation is fused by the compilerTo understand what Inductor did, take a look at the generated Triton kernels:
TORCH_LOGS="output_code" python implicit_fusion.py
The interesting part from the first kernel:
tmp0 = tl.load(in_ptr0 + ...) # load a chunk of input
tmp1 = tl_math.abs(tmp0) # apply abs in registers, inline
tmp2 = _tmp3 + tmp1 # accumulate into the running sum
Theabshappens in a register, right after the load, inside the reduction kernel. There’s no separateabskernel and no per-element intermediate buffer. A second kernel, running in one block, reads the 2,048 partial results and reduces them to one final scalar. You can learn more about this in thisPyTorch blog post.
In practice,torch\.compilefused the operation but split the reduction into two stages: one kernel reduces the input to a small set of per-block partial sums, and a second kernel adds those partial sums into a single number. In the manual version, we made a different choice: using atomics to accumulate into one global float. Both versions move one GiB through global memory instead of 3 GiB. The compiler picked a reasonable strategy but not the same strategy we would choose.
This is both good and bad. The compiler gives us a fused kernel for free, but we trade away some predictability. A different dtype, a different shape, or a small operation added in the middle can produce a completely different set of kernels. Sometimes you get one kernel. Sometimes you get four. Sometimes a fusion you were counting on quietly stops happening, without warning.
Explicit kernel fusion: Callbacks, iterators, and epilogshttps://developer.nvidia.com/blog/kernel-fusion-in-nvidia-cuda-optimizing-memory-traffic-and-launch-overhead/#explicit_kernel_fusion_callbacks_iterators_and_epilogs
What if we want the productivity of Python with the predictability of a hand-written kernel? That is whatcuda\.computeis for.
cuda\.computeexposes device-wide algorithms such asreduce,scan,sort, andtransform, along with a set of iterators that you can compose to control how those algorithms read their input. You do not write a kernel directly. Instead, you compose a computation and pass it to an algorithm, and the fusion happens because you explicitly asked for it.
The samesum\(abs\(x\)\)operation incuda\.computelooks like this:
import torch
import numpy as np
from cuda.compute import (
Determinism, OpKind, TransformIterator, reduce_into,
)
x = torch.randn(N, device="cuda")
out = torch.empty(1, dtype=torch.float32, device="cuda")
h_init = np.array([0.0], dtype=np.float32)
abs_it = TransformIterator(x, lambda a: abs(a))
reduce_into(
d_in=abs_it, d_out=out, num_items=N,
op=OpKind.PLUS, h_init=h_init,
determinism=Determinism.NOT_GUARANTEED,
)
TransformIteratordoesn’t allocate anything or launch a kernel. It’s a lazy view that says, “when someone asks me for elementi, readx\[i\]and applyabsto it.”reduce\_intothen performs a single-device-wide reduction, pulling elements through the iterator.
Theabsruns inline, in registers, as each element is loaded from global memory. There’s no per-element intermediate buffer at any point. Also note thatcuda\.computeworks with PyTorch Tensors. It supports inputs that implement DLPack or the CUDA Array API, such as CuPy arrays.
Once you separate “what the data is” from “how to walk it,” producing elementibecomes an arbitrary inline computation: a load, a transform, a gather, or a struct unpack. The algorithm threads these operations into a single pass with no intermediate storage. This kind of solution is deterministic, portable, and composable, making it ideal for library authors and performance applications.
If we profile this with Nsight Systems, we get something similar to a single fused kernel produced by a few lines of Python:
*Figure 5. Nsight Systems profile for the explicit fusion fromcuda\.compute. One kernel calledNondeterministicDeviceReduceAtomicKernel*TheCUB library, on whichcuda\.computeis built, produces a single fused kernel optimized for the operation we described. It provides battle-tested implementations for many algorithms, along with the guarantees CCCL provides, such as optimizations for new hardware.
Unlike a compiler approach, it won’t silently change the fusion strategy because of a small code change, a different compiler version, or a compiler decision, because you explicitly decide what should be fused. This gives you more deterministic results while still enabling control over the final behavior, including details likefloating-point determinism.
MetricNaiveManual Fusiontorch\.compile``cuda\.computeKernels Launched2121Per-element intermediate bufferyesnononoSpeedup vs. naive–3x3x3xWho decides what gets fused?–YouCompilerYouDeterminism–FullVariesPredictableTable 2. Summary of all kernel fusion methodsAll three fused methods end up in the same place with 1 GiB of memory traffic and about a 3x speedup over the naive version. With manual fusion, you write the kernel, so the fusion is explicit and under your control. With compiler fusion, the compiler chooses for you.
Withcuda\.compute, you write the composition using Python ergonomics. Like manual fusion, this provides an explicit guarantee that the transform is fused into the reduction. The added benefit is that this composition is backed by CUB, so you get a well-tested, highly optimized implementation under the hood instead of a custom kernel you have to maintain and tune yourself.
cuda\.computegives Python developers access to the same iterator-based, deterministic fusion that Thrust and CUB have long provided in C++. It enables higher-level APIs and frameworks to perform as if they had custom kernels, without requiring developers to write those kernels themselves.
You can learn more aboutcuda.computein docs andcuda-samplesfor Python.
相似文章
@PyTorch: https://bit.ly/4yawNqB..*
这篇来自PyTorch的博客文章介绍了针对LayerNorm和RMSNorm等归一化操作的新型内核融合技术,通过减少内存IO开销实现了显著的加速。这些技术包括Lazy Pre-Norm和Multi-CTA Norm Fusion,在与GEMM融合时可隐藏高达90%的归一化延迟,而FlashNormAttention算法可实现高达35%的内核加速。
一个可定制的编译器,用于为AI模型生成高效的融合GPU内核 [P]
作者介绍了一款用 Python 编写、高度可定制且易于修改的 ML 编译器。该编译器通过多级 IR 流水线将 LLMs 转换为优化的 CUDA 内核,在特定操作上实现了与 PyTorch 相当甚至更优的性能。文章详细阐述了该编译器的优化过程、降级规则以及用于生成高效融合 GPU 内核的 CLI 用法。
@PyTorch:一个运行时,多种GPU架构,零厂商特定模型代码。在这篇博文中,TokenSpeed团队 @l…
TokenSpeed-Kernel是一个可移植、高性能的内核系统,用于LLM推理,实现零厂商特定模型代码,并支持多种GPU架构,在AMD MI355X上实现高达3.6倍的吞吐量提升。
@h100envy: Meta的PyTorch核心工程师将CUDA内核编写变成了一项13分钟的运动——比1500美元的GPU编程课程还要好……
Meta的一位PyTorch核心工程师展示了一个快速的CUDA内核优化循环,其性能优于昂贵的训练营,获胜代码通过KernelBot竞赛合并到了PyTorch中。
使用CUDA内核重写模型推理:瓶颈不仅仅是GEMM [P]
作者描述了构建FlashRT的过程,这是一个以CUDA为核心的推理运行时,通过使用C++/CUDA内核重写模型推理路径,来解决小批量/实时工作负载中超出GEMM的瓶颈,在Jetson Thor和RTX 5090上实现了显著的延迟改进。文章讨论了关于精度的经验(FP8有帮助,FP4好坏参半)以及绕过通用运行时进行实时推理的必要性。