Fused Kernels

1What fusing means

Here is the unfused way to compute out = x*2 + 3 in torch:

tmp = x * 2.0          # kernel 1: read x, write tmp
out = tmp + 3.0        # kernel 2: read tmp, write out

Two kernels, four memory passes — even though the math is trivial. The GPU spends its time shipping data to and from HBM, not computing. The fused way:

x = tl.load(x_ptr + offs, mask=mask)
tl.store(out_ptr + offs, x * 2.0 + 3.0, mask=mask)   # one kernel, two passes

Fusion = read your inputs once, do the whole computation in registers, write once. Everything between the tl.load and the tl.store is free — register arithmetic doesn’t touch memory at all. That’s why an elementwise chain (add, multiply, exp, sigmoid — any sequence) costs the same as a single op once the data is loaded.

2The economics: counting memory passes

Every kernel launch costs two things: the launch overhead (~5–10 µs, which decode feels hard) and the memory passes. Count the passes for the unfused MLP below vs the fused one:

The fused kernel: read x once, read the three weight matrices once, write out once — 5 passes, and x is read once instead of twice. Nearly 3× less data movement, and 1 launch instead of 5.

3Fusing chains: the silu gate

Part A of the companion is the trivial chain (x*2 + 3). Part B is the real one — the core of SwiGLU, the activation that replaced GELU in modern LLMs:

@triton.jit
def silu_gate_kernel(g_ptr, u_ptr, out_ptr, N, BLOCK: tl.constexpr):
    pid = tl.program_id(0)
    offs = pid * BLOCK + tl.arange(0, BLOCK)
    mask = offs < N
    g = tl.load(g_ptr + offs, mask=mask)
    u = tl.load(u_ptr + offs, mask=mask)
    tl.store(out_ptr + offs, g * tl.sigmoid(g) * u, mask=mask)

Three elementwise ops (sigmoid, *, *) in one line, one pass: silu(g) * u. In torch that’s two kernels (or one fused F.silu and a multiply — still two passes over g and u). Here: load both, compute, store. This exact line — with different names — is what decode-lab’s swiglu kernel does between its matmuls.

4The fused MLP — decode-lab’s swiglu.py

The full MLP is three matmuls around that gate:

out = down(silu(x @ gate^T) * (x @ up^T))

Torch runs five kernels (table above). The fused kernel runs one: it computes gate/up tiles, applies silu in registers, and dots into the output — reusing the transpose-read trick from lesson 4 to read the weights sideways:

@triton.jit
def swiglu_mlp_kernel(x_ptr, gate_ptr, up_ptr, down_ptr, out_ptr, M, N_OUT, H, I,
                      N_ITERS_I: tl.constexpr, N_ITERS_H: tl.constexpr,
                      BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr,
                      BLOCK_H: tl.constexpr, BLOCK_I: tl.constexpr):
    pid_m = tl.program_id(0)
    pid_n = tl.program_id(1)

    x_bp = tl.make_block_ptr(x_ptr, (M, H), (H, 1), (pid_m * BLOCK_M, 0),
                             (BLOCK_M, BLOCK_H), (1, 0))
    out_bp = tl.make_block_ptr(out_ptr, (M, N_OUT), (N_OUT, 1),
                               (pid_m * BLOCK_M, pid_n * BLOCK_N),
                               (BLOCK_M, BLOCK_N), (1, 0))

    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
    for ki in range(0, N_ITERS_I):
        # transposed views (lesson 4 Part C): gate (I,H) read as (H,BLOCK_I), etc.
        gate_bp = tl.make_block_ptr(gate_ptr, (H, I), (1, H), (0, ki * BLOCK_I),
                                    (BLOCK_H, BLOCK_I), (0, 1))
        up_bp = tl.make_block_ptr(up_ptr, (H, I), (1, H), (0, ki * BLOCK_I),
                                  (BLOCK_H, BLOCK_I), (0, 1))
        down_bp = tl.make_block_ptr(down_ptr, (I, N_OUT), (1, I),
                                    (ki * BLOCK_I, pid_n * BLOCK_N),
                                    (BLOCK_I, BLOCK_N), (0, 1))

        g = tl.zeros((BLOCK_M, BLOCK_I), dtype=tl.float32)
        u = tl.zeros((BLOCK_M, BLOCK_I), dtype=tl.float32)
        x_win = x_bp                                   # fresh window per I-chunk
        for kh in range(0, N_ITERS_H):
            x_t = tl.load(x_win, boundary_check=(0, 1), padding_option="zero")
            g_t = tl.load(gate_bp, boundary_check=(0, 1), padding_option="zero")
            u_t = tl.load(up_bp, boundary_check=(0, 1), padding_option="zero")
            g = tl.dot(x_t, g_t, g)                    # g  += x @ gate^T chunk
            u = tl.dot(x_t, u_t, u)                    # u  += x @ up^T chunk
            x_win = tl.advance(x_win, (0, BLOCK_H))
            gate_bp = tl.advance(gate_bp, (BLOCK_H, 0))
            up_bp = tl.advance(up_bp, (BLOCK_H, 0))

        hidden = g * tl.sigmoid(g) * u                 # silu(gate) * up — fused
        d_t = tl.load(down_bp, boundary_check=(0, 1), padding_option="zero")
        acc = tl.dot(hidden, d_t, acc)                 # acc += hidden @ down^T chunk

    tl.store(out_bp, acc.to(out_ptr.dtype.element_ty), boundary_check=(0, 1))

5When NOT to fuse

Fusion is not free. Everything fused lives in registers, and registers are the scarcest resource on a GPU:

The engineering rule: fuse memory-bound chains; leave compute-bound kernels alone. The swiglu kernel’s nested loops exist precisely because fusing everything into one giant tile would blow past shared memory — so it chunks instead.

6The pass counter

Watch the difference with your own eyes — run the two versions of the MLP and count the memory passes:

7Homework — seriously, do this

Run lesson5_fused_kernels.py on Colab, then try each of these by modifying the code. Hints are collapsible — try it first, peek only if stuck.

8Your learning journey