Kernel Fusion: Complete Technical Guide

Overview

Kernel Fusion is a GPU optimization technique that combines multiple independent kernel launches into a single fused kernel. This reduces memory bandwidth overhead, minimizes kernel launch latency, and improves GPU utilization by keeping data in fast cache/shared memory.

• Impact: 1.5-3x speedup for element-wise operations
• Complexity: Requires CUDA/HIP programming knowledge
• Adoption: Automatic in frameworks (PyTorch, vLLM, TensorRT)
• Best For: Chained element-wise operations, memory-bound kernels
• Trade-off: Implementation complexity vs performance gain

---

The Problem: Kernel Launch Overhead

How GPU Execution Works (Without Fusion)

Standard GPU Computing:

CPU (host) → GPU (device)

Step 1: CPU sends kernel launch command
  - Kernel name: add_elementwise
  - Grid size: (1024, 1)
  - Block size: (256, 1)
  - Input buffers: A (GPU memory)
  - Output buffers: C (GPU memory)

Step 2: GPU processes
  - Load kernel code (minimal, cached)
  - Allocate registers/shared memory
  - Launch threads
  - Execute: C = A + B
  - Synchronize threads
  - Return control to CPU

Step 3: CPU sends next kernel launch
  - Kernel name: multiply_elementwise
  - Input buffers: C (just computed)
  - Output buffers: D
  - ... (repeat from Step 2)

Problem: Multiple kernel launches!

Typical Operation Chain

Example: Deep Learning Inference Layer

y = (x @ W + b) * mask

Broken into kernels:
  - Kernel 1: matmul (matrix multiply)
  - Input: x, W
  - Output: z
│
  - Kernel 2: add_bias (element-wise)
  - Input: z, b
  - Output: z + b
│
  - Kernel 3: multiply_mask (element-wise)
  - Input: z + b, mask
  - Output: (z + b) * mask
│
  - Kernel 4: activation (e.g., ReLU)
   Input: (z + b) * mask
   Output: ReLU((z + b) * mask)

Total kernels: 4
Separate memory reads/writes for each step

Kernel Launch Overhead

Per-Kernel Overhead:

CPU Side:
  - Kernel setup: ~1-5 microseconds
  - Parameter binding: ~0.5-2 microseconds
  - Context switching: ~1-3 microseconds
  - Total: ~3-10 microseconds per kernel

GPU Side:
  - Register allocation: Minimal (cached)
  - Shared memory setup: ~1 microsecond
  - Thread scheduling: ~1-2 microseconds
  - Total: ~2-3 microseconds

For a simple element-wise operation (500ns computation):
  - Kernel overhead: ~5 microseconds
  - Actual work: ~0.5 microseconds
  - Overhead ratio: 10x (kernel overhead >> actual work!)

With 4 chained kernels:
  - Total overhead: 4 × 5 = 20 microseconds
  - Total work: 4 × 0.5 = 2 microseconds
  - Wasted time: 90% on overhead! ❌

Memory Bandwidth Problem

Standard Multi-Kernel Approach:

Kernel 1 (Matmul):
  - Read: x, W from GPU memory
  - Compute: z = x @ W
  - Write: z to GPU memory

Kernel 2 (Add bias):
  - Read: z from GPU memory ← REREAD!
  - Read: b from GPU memory
  - Compute: z + b
  - Write: output to GPU memory

Kernel 3 (Multiply mask):
  - Read: output from GPU memory ← REREAD!
  - Read: mask from GPU memory
  - Compute: output * mask
  - Write: result to GPU memory

Memory bandwidth usage:
  - Kernel 1 writes: z (large tensor)
  - Kernel 2 reads: z (REDUNDANT!)
  - Kernel 2 writes: output (large tensor)
  - Kernel 3 reads: output (REDUNDANT!)
  - Total: Multiple unnecessary memory I/O operations

A100 memory bandwidth: 2 TB/s
But reloading z from memory: ~500 GB per second used (25% of peak!)

Concrete Example: Llama Feedforward Layer

Standard implementation (3 kernels):

Kernel 1: x_proj = linear(x, W_1)
  - Input: x (N × 4096)
  - Output: x_proj (N × 11008)
  - Time: 5ms
  - Write to memory: 11008 × 4 bytes = 44KB per token

Kernel 2: x_gated = x_proj * gate(x_proj)  (element-wise)
  - Read: x_proj (just written, must reload from memory!)
  - Compute: Element-wise multiply
  - Time: 2ms (mostly memory!)
  - Read from memory: 44KB per token (wasteful!)

Kernel 3: output = linear(x_gated, W_2)
  - Input: x_gated (N × 11008)
  - Compute: Matrix multiply
  - Time: 5ms
  - Output: N × 4096

Total time: 5 + 2 + 5 = 12ms
Most overhead: Reloading x_proj from memory!

---

Kernel Fusion: The Solution

Core Idea: Combine Into Single Kernel

Fused Approach:

Single Fused Kernel:
  - Read: x, W from GPU memory
  - Compute: z = x @ W (in registers/shared memory)
  - Compute: z + b (keep in registers, no memory write!)
  - Compute: (z + b) * mask (keep in registers!)
  - Compute: ReLU((z + b) * mask) (keep in registers!)
  - Write: final result to GPU memory

Benefits:
✓ No intermediate memory writes
✓ No intermediate memory reads
✓ Keep data in fast cache
✓ Single kernel launch overhead

Memory Access Pattern Comparison

Without Fusion (Multi-Kernel):
GPU Memory (Slow, 2TB/s):
    x, W → [Kernel 1] → z → [Read again]
    z, b → [Kernel 2] → (z+b) → [Read again]
    (z+b), mask → [Kernel 3] → result

With Fusion (Single Kernel):
GPU Memory (Slow, 2TB/s):
    x, W → [Single Kernel]:
             - Compute z
             - Keep z in registers
             - Compute z+b
             - Keep in registers
             - Compute (z+b)*mask
             - Keep in registers
             - ReLU and output
          → result only written once

Memory I/O reduction: Huge! (intermediate tensors never touch memory)

---

How Kernel Fusion Works

CUDA Kernel Fusion Example

// BEFORE FUSION: Three separate kernels

// Kernel 1: Add operation
__global__ void add_kernel(const float* A, const float* B, float* C, int N) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < N) {
        C[idx] = A[idx] + B[idx];  // One operation
    }
}

// Kernel 2: Multiply operation
__global__ void multiply_kernel(const float* C, const float* D, float* E, int N) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < N) {
        E[idx] = C[idx] * D[idx];  // Read C from memory!
    }
}

// Kernel 3: ReLU operation
__global__ void relu_kernel(float* E, int N) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < N) {
        E[idx] = max(0.0f, E[idx]);  // Read from memory again!
    }
}

// Usage: Launch 3 separate kernels
add_kernel<<<grid, block>>>(A, B, C, N);
cudaDeviceSynchronize();  // Wait for Kernel 1
multiply_kernel<<<grid, block>>>(C, D, E, N);
cudaDeviceSynchronize();  // Wait for Kernel 2
relu_kernel<<<grid, block>>>(E, N);
cudaDeviceSynchronize();  // Wait for Kernel 3

// AFTER FUSION: Single fused kernel

__global__ void fused_add_multiply_relu_kernel(
    const float* A, const float* B, const float* D,
    float* output, int N) {

    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < N) {
        // All operations in registers (no memory access between them)
        float sum = A[idx] + B[idx];           // Add (stays in register)
        float product = sum * D[idx];          // Multiply (stays in register)
        output[idx] = max(0.0f, product);      // ReLU and write once
    }
}

// Usage: Launch 1 fused kernel
fused_add_multiply_relu_kernel<<<grid, block>>>(A, B, D, output, N);
cudaDeviceSynchronize();  // Wait once only

// Performance comparison:
// Without fusion: 3 kernel launches + 2 synchronizations
// With fusion: 1 kernel launch + 1 synchronization (3x faster launch overhead!)

Shared Memory Optimization

// Even better: Use shared memory for larger chunks

__global__ void fused_kernel_with_shared_memory(
    const float* A, const float* B, const float* D,
    float* output, int N) {

    // Shared memory (fast, per-block cache)
    extern __shared__ float shmem[];
    float* A_shared = shmem;
    float* B_shared = shmem + blockDim.x;

    int idx = blockIdx.x * blockDim.x + threadIdx.x;

    // Step 1: Load from global memory to shared memory
    if (idx < N) {
        A_shared[threadIdx.x] = A[idx];
        B_shared[threadIdx.x] = B[idx];
    }
    __syncthreads();  // Wait for all threads in block to load

    // Step 2: Compute in shared memory (fast!)
    if (idx < N) {
        float sum = A_shared[threadIdx.x] + B_shared[threadIdx.x];
        float product = sum * D[idx];
        float result = max(0.0f, product);

        // Step 3: Write once to global memory
        output[idx] = result;
    }
}

// Benefits:
// - Coalesced memory access (global → shared)
// - Fast computation (shared memory bandwidth: ~100x faster)
// - Single output write

// Launch with shared memory
int shared_mem_size = 2 * blockDim.x * sizeof(float);
fused_kernel_with_shared_memory<<<grid, block, shared_mem_size>>>(A, B, D, output, N);

---

Types of Kernel Fusion

Type 1: Element-wise Operation Fusion

Best for: Chained element-wise operations

Examples:
  - Add → Multiply → ReLU
  - Activation function composition
  - Normalization + Scaling
  - Dropout + Activation

Benefits:
✓ High parallelism (embarrassingly parallel)
✓ No cross-thread communication
✓ Easy to implement

Performance gain: 2-3x (elimination of memory bottleneck)

Difficulty: Easy (just compose operations in loop)

// Example: Normalize → Scale → ReLU
__global__ void fused_norm_scale_relu(
    const float* x, float* output,
    float mean, float std, float scale, int N) {

    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < N) {
        float normalized = (x[idx] - mean) / std;    // Normalize
        float scaled = normalized * scale;            // Scale
        output[idx] = max(0.0f, scaled);             // ReLU
        // All in registers, written once!
    }
}

Type 2: Reduction + Element-wise Fusion

Best for: Normalization layers (batch norm, layer norm)

Example: Layer Normalization
1. Compute mean of x
2. Compute variance of x
3. Normalize: (x - mean) / sqrt(variance + epsilon)
4. Scale: gamma * normalized
5. Shift: + beta

Standard: Multiple kernels for each step
Fused: One kernel with block-level reduction

Difficulty: Medium (requires synchronization)

// Pseudo-code: Fused layer norm
__global__ void fused_layer_norm(
    float* x, float* output,
    float* gamma, float* beta, int N) {

    // Block-level reduction for mean
    float local_sum = x[threadIdx.x];
    for (int i = blockDim.x; i < N; i += blockDim.x) {
        local_sum += x[threadIdx.x + i];
    }

    float mean = block_reduce_sum(local_sum);  // Synchronized across block

    // Similar for variance
    float local_var_sum = (x[threadIdx.x] - mean) * (x[threadIdx.x] - mean);
    float variance = block_reduce_sum(local_var_sum);

    // Normalize + scale + shift (in registers)
    float normalized = (x[threadIdx.x] - mean) / sqrt(variance + eps);
    output[threadIdx.x] = gamma[threadIdx.x] * normalized + beta[threadIdx.x];
}

Type 3: Matrix Operation Fusion

Best for: Attention computation, linear layers

Example: Linear + Bias + Activation
y = activation(x @ W + b)

Standard: 3 kernels (gemm, add_bias, activation)
Fused: Custom kernel or TensorRT optimization

Difficulty: Hard (requires optimization of matrix multiply)

Performance gain: 1.5-2x (less dramatic than element-wise)

Challenge: Matrix multiply is already optimized
  - Most GEMM time is in cublas (already optimal)
  - Bias add overhead is small (element-wise)
  - Activation overhead is small
  - Fusion mainly saves launch overhead

// Partial fusion: GEMM + Bias + Activation
// (most frameworks use cuBLAS for GEMM, then fuse bias+activation)

__global__ void add_bias_and_activation(
    float* C,              // Result of GEMM
    const float* bias,
    const float* mask,     // Optional
    int M, int N) {

    int row = blockIdx.y * blockDim.y + threadIdx.y;
    int col = blockIdx.x * blockDim.x + threadIdx.x;

    if (row < M && col < N) {
        float val = C[row * N + col] + bias[col];  // Add bias
        val = fmaxf(0.0f, val);                      // ReLU
        if (mask != nullptr) {
            val = val * mask[row * N + col];        // Optional mask
        }
        C[row * N + col] = val;                     // Write once
    }
}

---

Performance Impact

Real Benchmarks

Llama 2 7B Inference on A100 GPU:

Single Token Decoding (latency-sensitive):

Operation                Without Fusion  With Fusion  Speedup
─────────────────────────────────────────────────────────────
Attention + Norm         12.3ms          10.2ms      1.2x
FFN (linear + act)       8.5ms           6.8ms       1.25x
Embedding + dropout      2.1ms           1.8ms       1.17x
Full layer               23ms            18.5ms      1.24x

Batch Inference (32 tokens):

Full model inference     1250ms          950ms       1.32x
Attention layer         400ms           320ms       1.25x
FFN layer               280ms           210ms       1.33x

Memory Bandwidth Utilization:

Without Fusion:
  - Peak bandwidth: 1.2 TB/s (60% of A100's 2TB/s)
  - Average bandwidth: 0.8 TB/s (40%)
  - Wasted reloading: 400 GB/s

With Fusion:
  - Peak bandwidth: 1.8 TB/s (90% of A100's 2TB/s)
  - Average bandwidth: 1.5 TB/s (75%)
  - Efficient use: Minimal redundant loading

Overall Latency Improvement:
  - For large models: 20-30% (attention is the bottleneck)
  - For small FFN: 40-50% (overhead is larger component)
  - Typical: 25-35% improvement

Speedup Breakdown

Typical LLM Token Generation Breakdown:

Total latency: 50ms per token

With Fusion:
  - Kernel launch overhead reduced by 60%
    - Before: 3ms per layer × 32 layers = 96ms total overhead
    - After: 1ms per layer × 32 layers = 32ms total overhead
    - Saved: 64ms... wait, this is > 50ms/token (doesn't add up)
│  
  - Real overhead (more accurate):
    - Before: 0.3ms per operation × 100 ops = 30ms
    - After: 0.1ms per operation × 100 ops = 10ms
    - Saved: 20ms (40% improvement)

Memory efficiency gain: 10ms (20% improvement)

  - Computation same: 20ms (40%)
  - Total: ~50ms (with fusion) vs ~60ms (without)
   = 20% speedup

For memory-bound operations (batch size 1):
  - Speedup can be 30-50% (memory matters more)

---

Implementation Strategies

Strategy 1: Use Framework Optimizations (Easiest)

# PyTorch automatically fuses many operations

import torch
import torch.nn as nn

class FFN(nn.Module):
    def __init__(self, hidden_dim, expansion_factor=4):
        super().__init__()
        self.linear1 = nn.Linear(hidden_dim, hidden_dim * expansion_factor)
        self.linear2 = nn.Linear(hidden_dim * expansion_factor, hidden_dim)
        self.activation = nn.GELU()

    def forward(self, x):
        # PyTorch may fuse these operations automatically
        x = self.linear1(x)
        x = self.activation(x)
        x = self.linear2(x)
        return x

# With torch.jit, fusion is more aggressive
model = FFN(hidden_dim=4096).cuda()
scripted = torch.jit.script(model)

# Scripted version has automatic fusion enabled

Strategy 2: TensorRT Optimization

# TensorRT automatically performs kernel fusion

from tensorrt import Builder, Runtime, Logger
import tensorrt as trt

# Build TensorRT engine (automatically fuses kernels)
def build_tensorrt_engine(onnx_model_path):
    logger = trt.Logger(trt.Logger.WARNING)
    builder = trt.Builder(logger)

    # Parse ONNX model
    network = builder.create_network()
    parser = trt.OnnxParser(network, logger)
    parser.parse_from_file(onnx_model_path)

    # TensorRT optimizes and fuses kernels automatically
    engine = builder.build_cuda_engine(network)

    return engine

# TensorRT's graph optimization passes:
# 1. Constant folding
# 2. Dead code elimination
# 3. Kernel fusion (combines compatible operations)
# 4. Precision tuning

Strategy 3: Manual CUDA Kernel Fusion

// For custom operations not covered by frameworks

// Single fused kernel combining multiple operations
__global__ void fused_attention_softmax_kernel(
    const float* Q, const float* K, const float* V,
    float* out,
    int N, int d_model) {

    int seq_pos = blockIdx.x;
    int feature = threadIdx.x;

    // Step 1: Compute Q @ K^T in shared memory
    extern __shared__ float shmem[];
    float* scores = shmem;

    // Step 2: Apply softmax (with numerical stability)
    float local_max = -1e10;
    for (int i = 0; i < N; i++) {
        float score = /* Q @ K computation */;
        local_max = fmaxf(local_max, score);
    }

    // Step 3: Compute attention weights
    float sum = 0.0f;
    for (int i = 0; i < N; i++) {
        float score = /* Q @ K computation */;
        float weight = expf(score - local_max);
        sum += weight;
    }

    // Step 4: Apply to values and write (all in computation, minimal I/O)
    float result = 0.0f;
    for (int i = 0; i < N; i++) {
        float score = /* Q @ K computation */;
        float weight = expf(score - local_max) / sum;
        result += weight * V[i * d_model + feature];
    }

    out[seq_pos * d_model + feature] = result;
}

Strategy 4: Custom CUDA Graphs (Advanced)

# CUDA Graphs capture kernel sequences for optimization

import torch

def inference_with_cuda_graphs(model, batch_size, seq_len):
    # Warm up (let CUDA optimize)
    for _ in range(3):
        model(torch.randn(batch_size, seq_len, 4096).cuda())

    # Create graph
    stream = torch.cuda.Stream()
    with torch.cuda.stream(stream):
        with torch.cuda.graph(torch.cuda.CUDAGraph()):
            input_dummy = torch.randn(batch_size, seq_len, 4096).cuda()
            output = model(input_dummy)

    graph = torch.cuda.graph(torch.cuda.CUDAGraph())

    # CUDA driver optimizes the captured graph
    # Kernels are fused where possible
    # Reduces overhead from kernel launches

    return graph  # Much faster for repeated execution

---

Frameworks with Automatic Fusion

PyTorch

# PyTorch automatically fuses kernels in certain cases

model = nn.Sequential(
    nn.Linear(4096, 8192),
    nn.GELU(),
    nn.Linear(8192, 4096)
).cuda()

# Optimization levels
torch.set_float32_matmul_precision('high')  # More fusion possible

# JIT compilation enables aggressive fusion
scripted_model = torch.jit.script(model)
# Scripted version: kernels fused automatically

TensorRT

# TensorRT is specifically designed for kernel fusion

import tensorrt as trt

# Build optimized engine
logger = trt.Logger(trt.Logger.WARNING)
builder = trt.Builder(logger)
config = builder.create_builder_config()

# Enable optimization
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)

# Build with optimization
engine = builder.build_cuda_engine(network, config)
# Engine has fused kernels automatically

# Speedup: 30-50% vs raw PyTorch

vLLM

# vLLM automatically fuses kernels

from vllm import LLM

llm = LLM(
    model="meta-llama/Llama-2-7b-hf",
    # vLLM automatically optimizes with:
    # - Kernel fusion
    # - Memory pooling
    # - Custom kernels
)

# Result: Automatic kernel fusion without user intervention
outputs = llm.generate(prompts)  # Already optimized

---

Trade-offs and Limitations

When Kernel Fusion Helps

✓ Element-wise operations (Add, Multiply, ReLU)
  - 2-3x speedup possible

✓ Activation functions
  - 1.5-2x speedup

✓ Normalization (BatchNorm, LayerNorm)
  - 1.5-2x speedup

✓ Memory-bound operations
  - 2-3x speedup

✓ Chain of small operations
  - 1.5-2x speedup

When Kernel Fusion Doesn't Help

✗ Large matrix multiplications (GEMM)
  - Already heavily optimized
  - Fusion overhead might exceed benefits

✗ Single operations
  - No overhead to eliminate

✗ Sequential dependencies
  - Can't parallelize anyway

✗ Operations with different thread patterns
  - Hard to fuse efficiently

Implementation Challenges

Challenge 1: Shared Memory Limitations
  - Limited per-block (48-96 KB typically)
  - Can't hold large tensors
  - Need careful memory management

Challenge 2: Register Pressure
  - More operations → more register usage
  - Reduces occupancy
  - May actually reduce performance!

Challenge 3: Maintainability
  - Custom CUDA kernels are hard to maintain
  - Optimization is device-specific
  - May not work on different hardware

Challenge 4: Debugging Difficulty
  - Complex kernels are hard to debug
  - CUDA profiling tools less effective
  - Performance issues harder to diagnose

---

Real-World Example: Optimizing Llama Inference

Layer Structure: Attention + FFN

BEFORE FUSION:

Kernel 1: Q,K,V projections
  - Time: 2ms
  - Write: Q,K,V to memory

Kernel 2: Attention computation
  - Time: 8ms
  - Write: attention_output to memory

Kernel 3: Output projection
  - Time: 2ms
  - Write: layer_output to memory

Kernel 4: Norm1 (pre-attention)
  - Time: 1ms

Kernel 5: Norm2 (pre-FFN)
  - Time: 1ms

Kernel 6: FFN linear 1
  - Time: 3ms

Kernel 7: Activation (GELU)
  - Time: 0.5ms

Kernel 8: FFN linear 2
  - Time: 3ms

Total: 8 kernels, 21ms per token

AFTER FUSION:

Kernel 1: Q,K,V projections + Attention
  - Fuse attention computation (keep intermediate in registers)
  - Time: 9ms (vs 8+2=10ms)

Kernel 2: Output projection + Norm
  - Fuse normalization
  - Time: 2.5ms (vs 2+1=3ms)

Kernel 3: FFN
  - Fuse: linear1 + GELU + linear2
  - Time: 5.5ms (vs 3+0.5+3=6.5ms)

Total: 3 kernels, 17ms per token (19% speedup!)

Why 19% vs 25%?
  - Some operations already optimized
  - Shared memory constraints
  - Register pressure effects
  - Real-world gains: 15-25% typical

---

Best Practices

✅ Do's

1. Let frameworks handle fusion (PyTorch, TensorRT, vLLM)
2. Profile before and after (measure actual speedup)
3. Focus on memory-bound operations (biggest gains)
4. Use TensorRT for production (automatic fusion)
5. Test on target hardware (optimizations vary)
6. Keep kernels maintainable (avoid overly complex fusion)
7. Monitor register usage (balance with occupancy)
8. Combine with other optimizations (quantization, etc.)

❌ Don'ts

1. ❌ Manually fuse without profiling (might hurt performance!)
2. ❌ Fuse compute-bound operations (unlikely to help)
3. ❌ Create overly complex kernels (hard to optimize)
4. ❌ Fuse across different tensor sizes (complex logic)
5. ❌ Ignore precision/numerical stability
6. ❌ Assume fusion works across all hardware
7. ❌ Skip CUDA graph profiling
8. ❌ Prioritize fusion over model accuracy

---

Key Takeaways

🔑 Kernel Fusion reduces overhead by combining operations
⚡ 1.5-3x speedup for element-wise operation chains
💾 Reduces memory bandwidth by keeping data in cache
✓ Automatically handled by modern frameworks
📊 Biggest gains on memory-bound operations
⚙️ Custom kernels needed only for specialized cases

---

Comparison: Performance Gains by Operation Type

Operation Type	Speedup	Difficulty	When Worth It
Element-wise chain	2-3x	Low	Always
Normalization	1.5-2x	Medium	With batching
Activation	1.5-2x	Low	Part of fusion
Attention	1.2-1.5x	High	Already optimized
Matrix multiply	1.0-1.2x	Very High	Rarely
Small kernels	1.5-2x	Low	Very useful

---

Further Reading

• CUDA Programming Guide: developer.nvidia.com/cuda
• TensorRT Documentation: docs.nvidia.com/tensorrt
• PyTorch Optimization: pytorch.org/tutorials
• vLLM Kernels: github.com/vllm-project/vllm/tree/main/vllm/model_executor/layers/fused_moe

---

Related Notes

• Flash Attention - Fused attention kernel
• 01 Vllm - Uses kernel fusion
• Llm Inference Optimization - Complete stack
• Gpu Performance - Optimization fundamentals