Writing Custom CUDA Kernels for Attention

The PyTorch Overhead

We were researching a highly customized sparse attention mechanism. While PyTorch is incredibly flexible, composing multiple small mathematical operations in standard PyTorch creates massive overhead. Each operation requires writing the intermediate tensor back to the GPU's slow High Bandwidth Memory (HBM) and reading it back for the next operation.

Triton to the Rescue

We spent three weeks writing a custom CUDA kernel using OpenAI's Triton language. Triton allows you to write Python-like code that compiles down to highly optimized PTX assembly, abstracting away the nightmare of manual thread block management in raw CUDA C++.

import triton
import triton.language as tl

@triton.jit
def fused_activation_kernel(x_ptr, y_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
    # Calculate thread pointers
    pid = tl.program_id(axis=0)
    block_start = pid * BLOCK_SIZE
    offsets = block_start + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements
    
    # Load from HBM to SRAM
    x = tl.load(x_ptr + offsets, mask=mask)
    
    # Perform math in ultra-fast SRAM (Fusion)
    y = tl.where(x > 0, x, 0.01 * x) # Leaky ReLU
    
    # Write back to HBM once
    tl.store(y_ptr + offsets, y, mask=mask)

By fusing the operations into a single kernel, we kept the intermediate tensors in the GPU's ultra-fast SRAM. Our training speed increased by 2.5x. If you are doing novel architectural research in 2026, you must learn Triton.