Fused Element-wise Kernels: Implementation and Performance Benefits

The "Why" of Kernel Fusion: Avoiding the Memory Wall

Imagine a typical sequence in a neural network layer: you perform a matrix multiplication, add a bias vector, and then apply an activation function like ReLU. Without fusion, this would translate into three separate kernel launches:

1. GEMM Kernel: Computes $Y=XW$, writing the result matrix $Y$ to global memory.
2. Bias-Add Kernel: Reads $Y$, adds bias $b$, and writes the result $Z=Y+b$ back to global memory.
3. ReLU Kernel: Reads $Z$, applies $max(0,Z)$, and writes the final result back to global memory.

The problem lies in the hand-offs between kernels. The intermediate tensors, $Y$ and $Z$, are written to and read back from slow GPU DRAM. This round-trip is the performance killer.

Kernel fusion collapses these steps into a single operation. Instead of separate kernels, we write one "mega-kernel" that performs all the operations in sequence. The intermediate results, like $Y+b$, never leave the GPU's fast on-chip memory; they are held in registers before the final result is written out.

The performance benefits are threefold:

1. Reduced Memory Traffic: This is the most significant gain. By eliminating writes and reads of intermediate tensors to global memory, we save immense memory bandwidth and latency.
2. Reduced Kernel Launch Overhead: Launching a kernel from the host CPU involves overhead. Fewer launches mean less time spent on this overhead.
3. Increased Arithmetic Intensity: As we learned in the context of the Roofline model, arithmetic intensity is the ratio of floating-point operations to bytes of memory accessed (FLOPs/Byte). By drastically reducing the bytes accessed for the same number of FLOPs, fusion increases a kernel's arithmetic intensity, allowing it to move from being memory-bound towards being compute-bound and thus achieve higher performance.

The "How" of Kernel Fusion: Programmer vs. Compiler

You might wonder if the nvcc compiler can automatically fuse your separate __global__ functions. As a general rule, it cannot. Kernel fusion is typically a manual refactoring process performed by the programmer. You, the developer, are responsible for identifying sequential, memory-bound kernels and rewriting them as a single, combined kernel.

However, once you've written this fused kernel, the compiler plays a crucial role. It optimizes your code by ensuring that the intermediate variables you create are stored in the fastest possible memory: registers.

Let's study a clear explanation of this process.

This manual approach is the foundation. However, the AI ecosystem has evolved to automate this. High-level frameworks like PyTorch have developed powerful JIT (Just-In-Time) compilers that can perform these fusions for you.

A Modern Approach: Fusion with torch.compile and Triton

The torch.compile feature in PyTorch is a prime example of automated fusion. It analyzes your Python code, identifies fusible patterns of operations, and generates highly optimized, fused kernels using a backend called Triton. Triton is a Python-based language created by OpenAI for writing efficient GPU kernels, which offers a higher level of abstraction than raw CUDA C++.

Let's watch a practical demonstration of this process. The speaker will take a Deep Learning Recommendation Model (DLRM), identify performance bottlenecks with the profiler, and then use torch.compile to automatically generate fused kernels, resulting in a dramatic speedup.

Implementing a Fused Bias-Add-ReLU Kernel

The video showed us what a fused kernel generated by torch.compile looks like. Now, let's examine the code for this pointwise_add_relu kernel in both Triton and CUDA C++, using the code provided in the blog post that accompanies the video.

Let's break down the logic of the CUDA kernel, as it's a perfect template for an element-wise fused operation:

__global__ void pointwise_add_relu_fusion_512(
    float* in_out_ptr0, // Input/Output tensor (e.g., shape [B, N])
    const float* in_ptr0,   // Bias tensor (e.g., shape [N])
    const int XBLOCK)       // Number of columns N
{
    const int xnumel = 65536; // Total elements B * N
    const int N = 512;        // This should match XBLOCK

    // Standard 1D global thread indexing
    const int idx = blockIdx.x * blockDim.x + threadIdx.x;

    if (idx < xnumel) {
        // 'idx' is the linear index into the [B, N] tensor
        int linear_index = idx;

        // 'bias_index' maps the linear index to the correct column for the bias
        // This effectively broadcasts the 1D bias vector across all rows
        int bias_index = idx % N;

        // Load input value
        float input_val = in_out_ptr0[linear_index];
        
        // Load corresponding bias value
        float bias_val = in_ptr0[bias_index];

        // --- FUSION HAPPENS HERE ---
        // The intermediate result of the addition is held in a register
        float intermediate_val = input_val + bias_val;
        
        // The ReLU is applied to the register value
        float final_val = max(0.0f, intermediate_val);
        // -------------------------

        // Store the final result once. No intermediate writes to global memory.
        in_out_ptr0[linear_index] = final_val;
    }
}

This simple kernel beautifully illustrates the power of fusion. The intermediate_val variable lives ephemerally in a register, completely avoiding a costly round-trip to global memory.

Quantifying the Benefits

The most direct benefit of eliminating intermediate tensors is a reduction in memory allocations and deallocations. The video provides a clear visual of this.

This improved memory behavior translates directly into speed. The blog post includes a table showing that the custom Triton kernel significantly outperforms the standard eager-mode PyTorch implementation, which executes operations one by one. This speedup is the direct result of avoiding the memory wall through fusion.

Conclusion

In this lesson, you have been introduced to kernel fusion, a cornerstone of high-performance GPU programming. You've moved from the highly-tuned, monolithic world of GEMM to the more common scenario of optimizing a sequence of smaller, memory-bound operations.

Key Takeaways:

• Kernel fusion combines sequential kernel operations into a single kernel to eliminate round-trips to global memory for intermediate results.
• The primary benefits are a drastic reduction in memory bandwidth consumption and the elimination of kernel launch overhead.
• This technique increases a kernel's arithmetic intensity, allowing it to better utilize the GPU's computational resources.
• While traditionally a manual code refactoring process in CUDA C++, modern tools like torch.compile can automate fusion, often by generating high-level Triton code.
• A fused bias-add-ReLU kernel serves as a canonical example, where the intermediate result of the addition is held in a register before the activation function is applied.

In our next lesson, we will tackle a more complex and critical AI kernel: Layer Normalization. Implementing it efficiently will require us to combine our new understanding of fusion with the parallel reduction techniques you mastered in Module 7, demonstrating how these optimization principles build upon one another.