Start your own AI-powered blog — freeGet started →

Custom CUDA Kernels for LLMs: From Theory to Production

Q: When should I write a custom CUDA kernel vs using a library?

Write custom kernels when: (1) you need to fuse multiple operations to eliminate HBM roundtrips, (2) the specific operation is not covered by existing libraries, (3) you're targeting a new GPU architecture with specific features. Use libraries (PyTorch, vLLM, TensorRT-LLM) when: the operations are standard matmul/attention/layernorm — these are already highly optimized.

Q: Triton vs CUDA C++: which should I learn first?

Triton. It gives 90-95% of the performance with 10% of the code. You can write production-quality fused kernels in a few hours. Learn CUDA C++ when you need: (1) maximum performance for critical kernels, (2) WMMA (Warp Matrix Multiply-Accumulate) instructions, (3) custom memory management, (4) features not exposed by Triton.

Q: How do I debug a kernel that produces wrong results?

Start with: (1) test with small sizes (N=2, M=2, K=2), (2) compare against a CPU reference implementation, (3) use torch.allclose with atol=1e-2 (Triton accumulates in FP32 while PyTorch may not), (4) check for integer overflow in index calculations, (5) verify mask handling for non-power-of-2 dimensions, (6) check for race conditions in shared memory writes.

Q: What's the best kernel optimization for LLM inference?

Fused attention with KV cache management. The attention mechanism dominates inference time for long contexts, and the KV cache size determines the maximum batch size and context length. Flash Attention kernels combined with page attention (vLLM-style) and GQA/MQA give the largest single improvement for production LLM inference.

Q: What GPU is best for custom kernel development in 2026?

H100/B200 for maximum performance (FP8 Tensor Cores, 3.35 TB/s HBM3, 132 SMs). RTX 5090 for development (same architecture, lower cost). The gap between consumer and datacenter GPUs is narrowing — Blackwell consumer GPUs support FP8 Tensor Cores and large L2 caches, making them viable for kernel development and testing.

Synor

May 25, 20268 min read

Podcast episode2 voices

0:36

Custom CUDA Kernels for LLMs: From Theory to Production

Detailed texture of yellow corn kernels, perfect for backgrounds. Photo by micka randrianjafisolo on Pexels

Quick Answer: Custom CUDA kernels are the primary reason LLM inference is 10-100x faster in 2026 than naive PyTorch implementations. The critical optimizations: kernel fusion (combine multiple operations into one kernel — eliminates global memory roundtrips, typically 2-5x speedup), Flash Attention (tiling attention computation to fit in SRAM instead of HBM — 2-4x for attention alone, enables 128K+ context), quantized matrix multiplication (FP8/INT8 tensor cores — 2x throughput over FP16), and specialized LLM kernels (GQA/MQA attention, fused MoE, speculative decoding kernels). The stack: Triton for productivity (write in Python-like DSL, compiles to PTX — 90% of custom kernel perf with 10% of the code), CUTLASS for template-based optimized GEMM, CUDA C++ for maximum control, and TensorRT-LLM/vLLM for pre-built production kernels. Key benchmark: a fused attention + MLP + residual kernel achieves 340 TFLOPS on H100 vs 120 TFLOPS for separate PyTorch operations — a 2.8x improvement.

GPU Memory Hierarchy and Why It Matters

The Memory Bottleneck

Every CUDA kernel optimization starts with understanding the GPU memory hierarchy:

code

Memory Hierarchy (H100 SXM):
┌──────────────────────────────────────────────────┐
│  HBM3 (High Bandwidth Memory)                    │
│  ~80 GB total, 3.35 TB/s bandwidth               │
│  ⚡ ~20-40 pJ/byte access energy                 │
├──────────────────────────────────────────────────┤
│                      ↑ 40x slower ↑              │
├──────────────────────────────────────────────────┤
│  L2 Cache (shared L2)                             │
│  50-60 MB, ~10 TB/s bandwidth                     │
├──────────────────────────────────────────────────┤
│                      ↑ 5-10x slower ↑            │
├──────────────────────────────────────────────────┤
│  Shared Memory / L1 (per SM)                      │
│  228 KB per SM (132 SM total = ~30 MB total)      │
│  ~100 TB/s aggregate bandwidth                    │
│  ⚡ ~2-5 pJ/byte access energy                   │
├──────────────────────────────────────────────────┤
│                      ↑ 10-20x slower ↑           │
├──────────────────────────────────────────────────┤
│  Registers (per thread)                           │
│  65536 32-bit registers per SM                    │
│  ~1 cycle access, ~1 pJ/byte                     │
└──────────────────────────────────────────────────┘

Key insight: Reading from HBM costs ~40x more energy
and is ~40x slower than reading from shared memory.

The "Roof" Model (Roofline Analysis)

code

Performance (TFLOPS)
  ▲
  │     ┌───────────────────────── Compute-bound
  │     │                          (not enough math)
  │     │
  │     │           ▲
  │     │           │ Peak compute: 2000 TFLOPS (FP8)
  │     │           │  989 TFLOPS (FP16)
  │  ───┤───────●───┴─── Ridge point
  │     │      ●
  │     │    ●          ▲ Arithmetic Intensity ▲
  │     │  ●   Memory-bound  │  (FLOPs/byte)
  │     │●  (limited by BW)  │
  │     ●
  │   ●
  │ ●
  ●───────┴──────────────────────────────▶
  Arithmetic Intensity (FLOPs/byte)

Goal: Move kernels from left (memory-bound) to right (compute-bound)
by increasing arithmetic intensity through kernel fusion.

Arithmetic Intensity Guide

Operation	Arithmetic Intensity (FP16)	Bound By
Element-wise (ReLU, add, multiply)	<1 FLOP/byte	Memory
LayerNorm	2-5 FLOPs/byte	Memory
Softmax	5-10 FLOPs/byte	Memory
GELU	10-20 FLOPs/byte	Memory
MatMul (small M=1, N=4096, K=4096)	50 FLOPs/byte	Memory
MatMul (large M=4096, N=4096, K=4096)	2000+ FLOPs/byte	Compute
Flash Attention (long context)	300+ FLOPs/byte	Compute
Fused kernel (MLP+Attention+Residual)	500+ FLOPs/byte	Compute

Kernel Fusion: The Single Most Important Optimization

The Fusion Problem

A standard transformer layer executes ~15-20 separate GPU kernel launches:

code

Naive PyTorch: 15 kernel launches per layer

1. Linear (QKV projection)       ← kernel launch
2. Reshape + permute             ← kernel launch  
3. Attention score (Q·K^T)       ← kernel launch
4. Softmax                       ← kernel launch
5. Attention output (score·V)    ← kernel launch
6. Linear (output projection)    ← kernel launch
7. Residual add                  ← kernel launch
8. LayerNorm                     ← kernel launch
9. Linear (FC1)                  ← kernel launch
10. GELU activation              ← kernel launch
11. Linear (FC2)                 ← kernel launch
12. Residual add                 ← kernel launch
13. LayerNorm                    ← kernel launch
14. Dropout                      ← kernel launch
15. RoPE                         ← kernel launch

Each launch: ~5-50μs overhead, plus HBM roundtrip per output

Fused Kernel Design

code

Fused kernel: 3 kernel launches per layer

Kernel 1: QKV projection + RoPE + reshape + permute
  → Reads: weights, input, position embeddings
  → Writes: Q, K, V tensors (fused computation)

Kernel 2: Flash Attention + output projection + residual
  → Reads: Q, K, V, weights, residual
  → Writes: attention output (no intermediate HBM writes)

Kernel 3: FC1 + GELU + FC2 + residual + LayerNorm
  → Reads: attention output, weights, residual
  → Writes: layer output (single HBM write)

Fusion Speedup

Transformer Layer	Kernel Launches	Memory Reads (GB)	Time (ms)	Speedup
Naive PyTorch	15	12.4	8.2	1.0x
Ops fused (5 kernels)	5	8.1	4.1	2.0x
Full fusion (3 kernels)	3	5.3	2.9	2.8x
Custom Flash attention	1-2	3.8	1.8	4.6x

Flash Attention: Tiled Attention in SRAM

The Standard Attention Problem

Standard attention requires computing and storing the full N×N attention matrix in HBM:

code

Standard Attention:
Memory: O(N²) for N×N attention matrix
  → N=4096: 16M elements = 64 MB (fits HBM but not SRAM)
  → N=32768: 1B elements = 4 GB (HBM only)
  → N=131072: 17B elements = 68 GB (exceeds most GPU memory)

Compute: O(N²d) where d is head dimension
  → N=4096, d=128: 2.1 TFLOPS per layer
  → N=32768: 137 TFLOPS per layer

Flash Attention Algorithm

code

Flash Attention (Forward Pass):
Given: Q, K, V ∈ ℝ^(N×d) stored in HBM
Target: O = softmax(QK^T/√d) · V

Algorithm:
1. Divide Q, K, V into blocks of size B that fit in SRAM
2. For each Q block (Q_i):
   For each KV block (K_j, V_j):
     a. Load Q_i, K_j, V_j from HBM → SRAM   (one block at a time)
     b. Compute S_ij = Q_i · K_j^T / √d       (in SRAM)
     c. Compute row-wise softmax of S_ij      (online softmax)
     d. Accumulate partial output O_i += softmax(S_ij) · V_j
     e. Keep running statistics (max, sum) for correction
3. Write final O_i to HBM

Block size B determined by SRAM:
  B ≤ SRAM_size / (3 * d * 2)  (Q, K, V blocks + output)
  For d=128, SRAM=128KB: B ≈ 170 tokens

Key trick: Online softmax allows incremental computation
without materializing the full S matrix.

Online Softmax

python

import torch
import triton
import triton.language as tl

@triton.jit
def flash_attention_kernel(
    q_ptr, k_ptr, v_ptr, o_ptr,
    N, d,
    stride_q, stride_k, stride_v, stride_o,
    BLOCK_SIZE: tl.constexpr,
):
    """Flash attention kernel with online softmax."""

    pid = tl.program_id(0)

    # Load Q block
    q_offset = pid * BLOCK_SIZE * stride_q
    q = tl.load(q_ptr + q_offset + tl.arange(0, BLOCK_SIZE)[:, None] * stride_q + tl.arange(0, d)[None, :])

    # Initialize online softmax statistics
    m_i = tl.full([BLOCK_SIZE], -float('inf'), dtype=tl.float32)
    l_i = tl.zeros([BLOCK_SIZE], dtype=tl.float32)
    acc = tl.zeros([BLOCK_SIZE, d], dtype=tl.float32)

    # Process K,V blocks
    for start_kv in range(0, N, BLOCK_SIZE):
        k = tl.load(k_ptr + start_kv * stride_k + tl.arange(0, BLOCK_SIZE)[:, None] * stride_k + tl.arange(0, d)[None, :])
        v = tl.load(v_ptr + start_kv * stride_v + tl.arange(0, BLOCK_SIZE)[:, None] * stride_v + tl.arange(0, d)[None, :])

        # Compute attention scores
        s = tl.dot(q, tl.trans(k)) / (d ** 0.5)  # B×B block

        # Online softmax
        m_new = tl.maximum(m_i, tl.max(s, axis=1))
        alpha = tl.exp(m_i - m_new)
        beta = tl.exp(s - m_new[:, None])

        # Update statistics
        l_new = alpha * l_i + tl.sum(beta, axis=1)

        # Accumulate output
        acc = acc * alpha[:, None] + tl.dot(beta.to(tl.float16), v.to(tl.float16))

        # Update running statistics
        m_i = m_new
        l_i = l_new

    # Normalize by sum of exponentials
    acc = acc / l_i[:, None]

    # Write output
    tl.store(o_ptr + pid * BLOCK_SIZE * stride_o + tl.arange(0, BLOCK_SIZE)[:, None] * stride_o + tl.arange(0, d)[None, :], acc)

Flash Attention Performance

Model	Context Length	Standard Attention (ms)	Flash Attention (ms)	Speedup
Llama 3 8B	8K	4.2	1.8	2.3x
Llama 3 8B	32K	32.1	5.4	5.9x
Llama 3 8B	128K	512+ (OOM)	28.3	18x+
Llama 3 70B	8K	18.7	8.2	2.3x
Llama 3 70B	32K	148	24.1	6.1x

Tile-Based Matrix Multiplication

The GEMM Problem

Matrix multiplication (GEMM) is the dominant operation in LLMs — 60-80% of total compute. The naive implementation is memory-bound. Tiling makes it compute-bound.

code

Naive MatMul:  C[M,N] = A[M,K] × B[K,N]

for m in range(M):          # Load A[m,:] from HBM
    for n in range(N):      # Load B[:,n] from HBM
        for k in range(K):  # Compute dot product
            C[m,n] += A[m,k] * B[k,n]

HBM accesses: M*K + K*N + M*N (each element read M or N times!)

Tiled MatMul with Shared Memory

code

Tiled MatMul with inner product blocking:

For each tile (tm, tn):
  1. Load A[tm, :] tile into shared memory
  2. Load B[:, tn] tile into shared memory  
  3. Compute tile product in registers
  4. Accumulate into C[tm, tn]

HBM accesses: M*K + K*N (each element read once!)
                      ┌───────────────────────┐
                      │    Shared Memory        │
                      │  A_tile    B_tile      │
                      │  [BM×BK]  [BK×BN]      │
                      └───────┬───────────┬────┘
                              │           │
                    Load      ▼           ▼ Load
                   ┌──────────────────────────┐
                   │       HBM                 │
                   │  A[M×K]    B[K×N]        │
                   └──────────────────────────┘

Tiling optimization: Choose BM, BK, BN to maximize SRAM usage
and minimize HBM reads:
  BM*BK + BK*BN = SRAM_size (typically 128KB)

Tensor Core Utilization

python

@triton.jit
def matmul_kernel(
    a_ptr, b_ptr, c_ptr,
    M, N, K,
    stride_a, stride_b, stride_c,
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
    BLOCK_SIZE_K: tl.constexpr,
    GROUP_SIZE_M: tl.constexpr,
):
    """Tiled matrix multiplication using Tensor Cores."""

    pid = tl.program_id(0)
    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)

    # Group SMs for better L2 cache utilization
    group_id = pid // num_pid_n
    first_pid_m = group_id * GROUP_SIZE_M
    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
    pid_m = first_pid_m + ((pid % num_pid_n) * group_size_m) // num_pid_n
    pid_n = (pid % num_pid_n) * group_size_m % num_pid_n

    offs_am = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_bn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    offs_k = tl.arange(0, BLOCK_SIZE_K)

    a_ptrs = a_ptr + offs_am[:, None] * stride_a + offs_k[None, :] * 1
    b_ptrs = b_ptr + offs_k[:, None] * stride_b + offs_bn[None, :] * 1

    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)

    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
        a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0)
        b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0)

        # Tensor Core matmul
        accumulator = tl.dot(a, b, acc=accumulator)

        a_ptrs += BLOCK_SIZE_K * 1
        b_ptrs += BLOCK_SIZE_K * stride_b

    c = accumulator.to(tl.float16)

    offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    c_ptrs = c_ptr + offs_cm[:, None] * stride_c + offs_cn[None, :] * 1
    c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N)

    tl.store(c_ptrs, c, mask=c_mask)

Fused Activation and Normalization Kernels

Fused MLP Block

The standard MLP block in transformers executes: Linear → Activation → Linear. These three operations can be fused into one kernel.

python

@triton.jit
def fused_mlp_kernel(
    x_ptr, w1_ptr, w2_ptr, out_ptr,
    M, N, K,  # M: tokens, N: hidden, K: intermediate
    stride_x, stride_w1, stride_w2, stride_out,
    BLOCK_SIZE: tl.constexpr,
):
    """Fused MLP: x → Linear1 → GELU → Linear2 → output."""

    pid = tl.program_id(0)
    row_start = pid * BLOCK_SIZE

    # Load input block
    x = tl.load(x_ptr + row_start * stride_x + tl.arange(0, BLOCK_SIZE)[:, None] * stride_x + tl.arange(0, N)[None, :])

    # Reduce precision for intermediate compute
    x_fp16 = x.to(tl.float16)

    # First linear: hidden → intermediate (N → K)
    # Load w1 row by row (in blocks)
    hidden = tl.zeros([BLOCK_SIZE, K], dtype=tl.float32)
    for n in range(0, N, BLOCK_SIZE):
        w1_block = tl.load(w1_ptr + n * stride_w1 + tl.arange(0, BLOCK_SIZE)[:, None] * stride_w1 + tl.arange(0, K)[None, :])
        x_block = x_fp16[:, n:n + BLOCK_SIZE]
        hidden += tl.dot(x_block, w1_block)

    # GELU activation (fused)
    # GELU(x) ≈ x * 0.5 * (1 + tanh(√(2/π) * (x + 0.044715 * x^3)))
    hidden_fp16 = hidden.to(tl.float16)
    tanh_arg = hidden_fp16 * 0.79788456 * (1 + 0.044715 * hidden_fp16 * hidden_fp16)
    gelu_out = hidden_fp16 * 0.5 * (1 + tl.math.tanh(tanh_arg))

    # Second linear: intermediate → hidden (K → N)
    output = tl.zeros([BLOCK_SIZE, N], dtype=tl.float32)
    for k in range(0, K, BLOCK_SIZE):
        w2_block = tl.load(w2_ptr + k * stride_w2 + tl.arange(0, BLOCK_SIZE)[:, None] * stride_w2 + tl.arange(0, N)[None, :])
        gelu_block = gelu_out[:, k:k + BLOCK_SIZE]
        output += tl.dot(gelu_block, w2_block)

    # Write combined output (no intermediate HBM writes)
    tl.store(out_ptr + row_start * stride_out + tl.arange(0, BLOCK_SIZE)[:, None] * stride_out + tl.arange(0, N)[None, :], output.to(tl.float16))

Fused LayerNorm

python

@triton.jit
def fused_layernorm_kernel(
    x_ptr, gamma_ptr, beta_ptr, out_ptr,
    N, eps,
    stride_x, stride_out,
    BLOCK_SIZE: tl.constexpr,
):
    """Fused LayerNorm with mean/variance computation."""

    row = tl.program_id(0)
    x_ptrs = x_ptr + row * stride_x + tl.arange(0, N)

    # Load input
    x = tl.load(x_ptrs, mask=tl.arange(0, N) < N)

    # Compute mean
    mean = tl.sum(x, axis=0) / N

    # Compute variance (fused with mean subtraction)
    x_centered = x - mean
    variance = tl.sum(x_centered * x_centered, axis=0) / N

    # Normalize
    inv_std = 1.0 / tl.sqrt(variance + eps)
    x_norm = x_centered * inv_std

    # Scale and shift (gamma, beta)
    gamma = tl.load(gamma_ptr + tl.arange(0, N), mask=tl.arange(0, N) < N)
    beta = tl.load(beta_ptr + tl.arange(0, N), mask=tl.arange(0, N) < N)
    out = x_norm * gamma + beta

    tl.store(out_ptr + row * stride_out + tl.arange(0, N), out, mask=tl.arange(0, N) < N)

Creative display of the word 'OPTIMIZE' on a pink textured surface. Photo by Ann H on Pexels

GQA and MQA Attention Kernels

Why GQA/MQA Exist

Grouped Query Attention (GQA) and Multi-Query Attention (MQA) reduce the KV cache size and memory bandwidth requirements for inference:

code

Standard Multi-Head Attention (MHA):
  Q: 32 heads × 128d = 32 queries
  K: 32 heads × 128d = 32 key/value pairs
  KV cache per token: 32 × 128 × 2 × 2 bytes = 16 KB/token

Grouped Query Attention (GQA, 8 groups):
  Q: 32 heads × 128d = 32 queries (shared among groups)
  K: 8 groups × 128d = 8 key/value sets
  KV cache per token: 8 × 128 × 2 × 2 bytes = 4 KB/token
  → 75% reduction in KV cache memory!

Multi-Query Attention (MQA):
  Q: 32 heads × 128d = 32 queries
  K: 1 head × 128d = 1 key/value set
  KV cache per token: 1 × 128 × 2 × 2 bytes = 0.5 KB/token
  → 97% reduction vs MHA!

Custom GQA Kernel

python

@triton.jit
def gqa_attention_kernel(
    q_ptr, k_ptr, v_ptr, o_ptr,
    N, d, num_heads, num_kv_heads,
    stride_q, stride_k, stride_v, stride_o,
    BLOCK_HEADDIM: tl.constexpr,
    BLOCK_N: tl.constexpr,
):
    """Grouped Query Attention with fused KV heads."""

    # Each program processes one query head
    # Multiple query heads share the same KV head
    pid = tl.program_id(0)
    head_id = pid
    kv_head_id = head_id * num_kv_heads // num_heads

    # Load Q block for this head
    q = tl.load(q_ptr + head_id * N * BLOCK_HEADDIM + tl.arange(0, BLOCK_N)[:, None] * N + tl.arange(0, BLOCK_HEADDIM)[None, :])

    # Initialize accumulator
    acc = tl.zeros([BLOCK_N, BLOCK_HEADDIM], dtype=tl.float32)
    m_i = tl.full([BLOCK_N], -float('inf'), dtype=tl.float32)
    l_i = tl.zeros([BLOCK_N], dtype=tl.float32)

    # Process KV in blocks
    for start in range(0, N, BLOCK_N):
        k = tl.load(k_ptr + kv_head_id * N * BLOCK_HEADDIM + start * N + tl.arange(0, BLOCK_N)[:, None] * N + tl.arange(0, BLOCK_HEADDIM)[None, :])
        v = tl.load(v_ptr + kv_head_id * N * BLOCK_HEADDIM + start * N + tl.arange(0, BLOCK_N)[:, None] * N + tl.arange(0, BLOCK_HEADDIM)[None, :])

        # Attention scores
        s = tl.dot(q, tl.trans(k)) / (BLOCK_HEADDIM ** 0.5)

        # Online softmax
        m_new = tl.maximum(m_i, tl.max(s, axis=1))
        alpha = tl.exp(m_i - m_new)
        beta = tl.exp(s - m_new[:, None])
        l_new = alpha * l_i + tl.sum(beta, axis=1)
        acc = acc * alpha[:, None] + tl.dot(beta.to(tl.float16), v.to(tl.float16))

        m_i = m_new
        l_i = l_new

    acc = acc / l_i[:, None]
    tl.store(o_ptr + head_id * N * BLOCK_HEADDIM + tl.arange(0, BLOCK_N)[:, None] * N + tl.arange(0, BLOCK_HEADDIM)[None, :], acc)

Quantized LLM Kernels (FP8, INT4, INT8)

Quantization-Aware Kernel Design

Quantized kernels reduce memory bandwidth and increase throughput at the cost of precision:

Precision	Memory per Param	Relative Throughput	Quality Impact
FP32	4 bytes	0.3x	None (baseline)
FP16/BF16	2 bytes	1.0x	None (baseline)
FP8 (E4M3/E5M2)	1 byte	1.8-2.0x	Negligible
INT8	1 byte	2.0x	Negligible (with smoothing)
INT4	0.5 bytes	2.5-3.0x	Small (with AWQ/GPTQ)
INT2/NF2	0.25 bytes	3.0-3.5x	Moderate (with calibration)

FP8 Kernel Using Tensor Cores

python

@triton.jit
def fp8_matmul_kernel(
    a_ptr, b_ptr, c_ptr,
    a_scale_ptr, b_scale_ptr,
    M, N, K,
    stride_a, stride_b, stride_c,
    BLOCK_SIZE_M: tl.constexpr,
    BLOCK_SIZE_N: tl.constexpr,
    BLOCK_SIZE_K: tl.constexpr,
):
    """FP8 matrix multiplication with per-tensor scaling."""

    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)

    offs_am = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_bn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    offs_k = tl.arange(0, BLOCK_SIZE_K)

    a_ptrs = a_ptr + offs_am[:, None] * stride_a + offs_k[None, :]
    b_ptrs = b_ptr + offs_k[:, None] * stride_b + offs_bn[None, :]

    # FP8 accumulation requires FP32 accumulator
    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)

    for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)):
        a_fp8 = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K)
        b_fp8 = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K)

        # Load FP8 per-block scaling factors
        a_scale = tl.load(a_scale_ptr + k)
        b_scale = tl.load(b_scale_ptr + k)

        # Dequantize to FP16 for Tensor Core matmul
        a_fp16 = a_fp8.to(tl.float16) * a_scale
        b_fp16 = b_fp8.to(tl.float16) * b_scale

        # Tensor Core matmul in FP16, accumulate in FP32
        accumulator = tl.dot(a_fp16, b_fp16, acc=accumulator)

        a_ptrs += BLOCK_SIZE_K
        b_ptrs += BLOCK_SIZE_K * stride_b

    # Store result
    c = accumulator.to(tl.float16)
    offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
    offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)
    c_ptrs = c_ptr + offs_cm[:, None] * stride_c + offs_cn[None, :]
    tl.store(c_ptrs, c, mask=(offs_cm[:, None] < M) & (offs_cn[None, :] < N))

Writing CUDA Kernels with Triton

Triton vs Raw CUDA

Feature	Triton	Raw CUDA C++	CUTLASS
Learning curve	Low (Python)	High (C++, PTX)	Medium (templates)
Performance (relative)	90-95%	100%	95-100%
Development time	Minutes	Days	Hours
Auto-tuning	Built-in	Manual	Built-in
Tiling automation	Automatic	Manual	Template-based
Tensor Core support	Yes	Yes	Yes
Debugging	Limited	Full (Nsight, cuda-gdb)	Moderate

Triton Auto-tuning

python

@triton.autotune(
    configs=[
        triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 64}, num_stages=4, num_warps=8),
        triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32}, num_stages=4, num_warps=4),
        triton.Config({'BLOCK_SIZE_M': 128, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 64}, num_stages=3, num_warps=8),
        triton.Config({'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 64, 'BLOCK_SIZE_K': 64}, num_stages=3, num_warps=8),
        triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 32}, num_stages=4, num_warps=4),
    ],
    key=['M', 'N', 'K'],
)
@triton.jit
def autotuned_matmul(a, b, c, M, N, K, **meta):
    """Auto-tuned matrix multiplication."""
    # Kernel implementation (same as tiled matmul above)
    ...

Production Triton Workflow

python

# Step 1: Profile existing PyTorch operation
import torch.utils.benchmark as benchmark

def profile_ops():
    x = torch.randn(4096, 4096, device='cuda', dtype=torch.float16)
    w = torch.randn(4096, 4096, device='cuda', dtype=torch.float16)

    # PyTorch matmul
    t_pytorch = benchmark.Timer(
        stmt='torch.matmul(x, w.T)',
        globals={'x': x, 'w': w}
    ).blocked_autorange()

    # Custom Triton matmul
    t_triton = benchmark.Timer(
        stmt='triton_matmul(x, w.T)',
        globals={'x': x, 'w': w, 'triton_matmul': custom_matmul}
    ).blocked_autorange()

    print(f"PyTorch: {t_pytorch.median * 1e6:.2f} μs")
    print(f"Triton:  {t_triton.median * 1e6:.2f} μs")
    print(f"Speedup: {t_pytorch.median / t_triton.median:.2f}x")

Production Kernel Optimization with CUTLASS

CUTLASS Architecture

CUTLASS provides template-based GPU kernels with automatic tile scheduling:

cpp

// CUTLASS 3.x GEMM example (H100 Tensor Core)
using Gemm = cutlass::gemm::device::Gemm<
    cutlass::half_t,                    // Input A type
    cutlass::layout::RowMajor,          // Layout A
    cutlass::half_t,                    // Input B type  
    cutlass::layout::ColumnMajor,       // Layout B
    cutlass::half_t,                    // Output type
    cutlass::layout::RowMajor,          // Layout output
    cutlass::half_t,                    // Accumulator type
    cutlass::arch::OpClassTensorOp,     // Use Tensor Cores
    cutlass::arch::Sm90,                // H100 (Hopper)
    cutlass::gemm::GemmShape<128, 128, 64>,  // Threadblock tile
    cutlass::gemm::GemmShape<64, 64, 64>,    // Warp tile
    cutlass::gemm::GemmShape<16, 8, 16>      // MMA tile
>;

Gemm gemm;
cutlass::Status status = gemm({
    {M, N, K},
    {a, a_stride},
    {b, b_stride},
    {c, c_stride},
    {d, d_stride},
    {alpha, beta}
});

CUTLASS vs Triton vs cuBLAS Performance

Operation	Problem Size	cuBLAS (TFLOPS)	CUTLASS (TFLOPS)	Triton (TFLOPS)
FP16 MatMul	8192×8192	450	465	440
FP16 MatMul	4096×4096	420	435	410
FP8 MatMul	8192×8192	820	850	800
INT8 MatMul	8192×8192	780	810	750
LayerNorm	Token 4096	N/A (PyTorch)	~95% HBM BW	~90% HBM BW

Benchmarking and Profiling Custom Kernels

Profiling Workflow

python

class CUDAKernelProfiler:
    """Production profiling framework for custom kernels."""

    def __init__(self):
        self.results = {}

    def profile_kernel(self, name, fn, inputs, warmup=10, iterations=100):
        """Profile kernel with warmup and multiple iterations."""

        # Warmup
        for _ in range(warmup):
            fn(**inputs)
        torch.cuda.synchronize()

        # Timed runs
        start_event = torch.cuda.Event(enable_timing=True)
        end_event = torch.cuda.Event(enable_timing=True)

        start_event.record()
        for _ in range(iterations):
            fn(**inputs)
        end_event.record()
        torch.cuda.synchronize()

        elapsed_ms = start_event.elapsed_time(end_event) / iterations

        # Calculate TFLOPS
        # Assume operation is matmul: 2*M*N*K FLOPs
        if 'M' in inputs and 'N' in inputs and 'K' in inputs:
            flops = 2 * inputs['M'] * inputs['N'] * inputs['K']
            tflops = flops / (elapsed_ms * 1e-3) / 1e12

        # Calculate memory bandwidth
        # Assume all operands read/written
        mem_bytes = sum(
            inp.numel() * inp.element_size() 
            for inp in inputs.values() 
            if isinstance(inp, torch.Tensor)
        )
        bw_gbs = mem_bytes / (elapsed_ms * 1e-3) / 1e9

        self.results[name] = {
            'time_ms': elapsed_ms,
            'tflops': tflops,
            'bandwidth_gbs': bw_gbs,
            'roof_bound': 'compute' if flops / mem_bytes > 100 else 'memory',
        }

        return self.results[name]

    def report(self):
        """Print formatted report."""
        print(f"{'Kernel':<30} {'Time (ms)':<12} {'TFLOPS':<12} {'BW (GB/s)':<12} {'Bound By':<12}")
        print("-" * 78)
        for name, metrics in self.results.items():
            print(f"{name:<30} {metrics['time_ms']:<12.3f} {metrics['tflops']:<12.1f} {metrics['bandwidth_gbs']:<12.1f} {metrics['roof_bound']:<12}")

Nsight Compute Metrics

bash

# Command-line profiling
ncu --set full -o kernel_profile python run_model.py

# Key metrics to check:
# 1. Achieved Occupancy (target: >50%)
# 2. SM Throughput (target: >80% peak)
# 3. HBM Read/Write Throughput (target: >70% peak for memory-bound)
# 4. L1 Hit Rate (target: >80%)
# 5. L2 Hit Rate (target: >60%)
# 6. Tensor Core Utilization (target: >70% for compute-bound)
# 7. Instruction Replay Overhead (target: <5%)

Common Kernel Bottlenecks and Fixes

Symptom	Root Cause	Fix
Low occupancy (<30%)	Too many registers per thread	Reduce BLOCK_SIZE, use `--maxrregcount`
Low SM throughput	Thread divergence	Restructure warp-level code, avoid branch
Low HBM bandwidth	Poor memory coalescing	Ensure contiguous memory access patterns
Low L1/L2 hit rate	Random access patterns	Use shared memory tiling, prefetching
Tensor Core underutilization	Wrong data layout	Pad dimensions to tensor core alignment
High instruction replay	Bank conflicts	Use proper shared memory padding

Key Takeaways

Kernel fusion is the single most impactful optimization for LLM inference, reducing 15+ PyTorch kernel launches per transformer layer to just 3 fused kernels—cutting memory reads by 57% and delivering 2.8x speedup by eliminating HBM roundtrips.
Flash Attention’s tiling algorithm (processing Q/K/V in SRAM-sized blocks) enables 128K+ context lengths with 18x+ speedup over standard attention by avoiding O(N²) HBM writes, using online softmax to compute attention incrementally without materializing full matrices.
Quantized matrix multiplication (FP8/INT8) doubles throughput over FP16 by halving memory bandwidth needs, with FP8 Tensor Core kernels achieving 1.8-2.0x speedup and negligible quality loss when paired with modern quantization techniques like AWQ.
GQA/MQA attention kernels reduce KV cache memory by 75-97% (e.g., 4KB/token for GQA vs 16KB/token for MHA) while maintaining accuracy, with custom Triton kernels fusing query/key/value operations to minimize memory traffic.
Tiled GEMM with shared memory (e.g., BM*BK + BK*BN ≤ 128KB SRAM) transforms memory-bound matrix multiplies into compute-bound operations, achieving 340 TFLOPS on H100 vs 120 TFLOPS for naive PyTorch implementations.
Fused activation/normalization kernels (e.g., MLP+GELU+LayerNorm) eliminate intermediate HBM writes by combining linear layers, activations, and residuals into single kernels, using register-level accumulation to maximize arithmetic intensity.

Frequently Asked Questions

When should I write a custom CUDA kernel vs using a library?

Write custom kernels when: (1) you need to fuse multiple operations to eliminate HBM roundtrips, (2) the specific operation is not covered by existing libraries, (3) you're targeting a new GPU architecture with specific features. Use libraries (PyTorch, vLLM, TensorRT-LLM) when: the operations are standard matmul/attention/layernorm — these are already highly optimized.

Triton vs CUDA C++: which should I learn first?

Triton. It gives 90-95% of the performance with 10% of the code. You can write production-quality fused kernels in a few hours. Learn CUDA C++ when you need: (1) maximum performance for critical kernels, (2) WMMA (Warp Matrix Multiply-Accumulate) instructions, (3) custom memory management, (4) features not exposed by Triton.

How do I debug a kernel that produces wrong results?

Start with: (1) test with small sizes (N=2, M=2, K=2), (2) compare against a CPU reference implementation, (3) use torch.allclose with atol=1e-2 (Triton accumulates in FP32 while PyTorch may not), (4) check for integer overflow in index calculations, (5) verify mask handling for non-power-of-2 dimensions, (6) check for race conditions in shared memory writes.

What's the best kernel optimization for LLM inference?

Fused attention with KV cache management. The attention mechanism dominates inference time for long contexts, and the KV cache size determines the maximum batch size and context length. Flash Attention kernels combined with page attention (vLLM-style) and GQA/MQA give the largest single improvement for production LLM inference.

What GPU is best for custom kernel development in 2026?

H100/B200 for maximum performance (FP8 Tensor Cores, 3.35 TB/s HBM3, 132 SMs). RTX 5090 for development (same architecture, lower cost). The gap between consumer and datacenter GPUs is narrowing — Blackwell consumer GPUs support FP8 Tensor Cores and large L2 caches, making them viable for kernel development and testing.

#Cuda Kernels #LLM Optimization #Flash Attention #GPU Kernel Fusion #Triton #Cutlass

Synor

1 followers

Deep dives on GPUs, decentralized AI, crypto, and open-source ML — buying guides, benchmarks, and tax/compliance explainers.

Comments

No comments yet. Be the first to share your thoughts!

Custom CUDA Kernels for LLMs: From Theory to Production

GPU Memory Hierarchy and Why It Matters

The Memory Bottleneck

The "Roof" Model (Roofline Analysis)

Arithmetic Intensity Guide

Kernel Fusion: The Single Most Important Optimization

The Fusion Problem

Fused Kernel Design

Fusion Speedup

Flash Attention: Tiled Attention in SRAM

The Standard Attention Problem

Flash Attention Algorithm

Online Softmax

Flash Attention Performance

Tile-Based Matrix Multiplication

The GEMM Problem

Tiled MatMul with Shared Memory

Tensor Core Utilization

Fused Activation and Normalization Kernels

Fused MLP Block

Fused LayerNorm

GQA and MQA Attention Kernels

Why GQA/MQA Exist

Custom GQA Kernel

Quantized LLM Kernels (FP8, INT4, INT8)

Quantization-Aware Kernel Design

FP8 Kernel Using Tensor Cores

Writing CUDA Kernels with Triton

Triton vs Raw CUDA

Triton Auto-tuning

Production Triton Workflow

Production Kernel Optimization with CUTLASS

CUTLASS Architecture

CUTLASS vs Triton vs cuBLAS Performance

Benchmarking and Profiling Custom Kernels

Profiling Workflow

Nsight Compute Metrics

Common Kernel Bottlenecks and Fixes

Related Reads

Key Takeaways

Frequently Asked Questions

When should I write a custom CUDA kernel vs using a library?

Triton vs CUDA C++: which should I learn first?

How do I debug a kernel that produces wrong results?

What's the best kernel optimization for LLM inference?

What GPU is best for custom kernel development in 2026?

Comments

More from Synor

FP8 vs INT4 Quantization: Production Quality vs Cost 2026

Bittensor Subnet Mining: Hardware Requirements and ROI

io.net Spot Instance Reliability: Availability Data 2026

Recommended for you

Flutter Custom Animations: From Basic to Production-Grade

How to Get a Free Professional Email Address (Custom Domain)

Welcome Email Series: Templates and Examples (2026)

Email Marketing vs Social Media: Which Has Better ROI in 2026?