Fused Kernels on Tensor Cores

The arithmetic

Why the bottom rung costs so much

Take a linear layer: GEMM, then add a bias, then GELU. Written as three kernels on an M×N output, that output crosses HBM five times — written once by the GEMM, read and written by bias, read and written by GELU. Fused, it crosses once.

The GEMM is compute-bound, so it barely notices. Bias and GELU do one flop per element, which puts them so far left on the roofline that they run at pure memory speed. At 4096×4096 in FP16 that is roughly 130 MB of traffic buying you 34 million flops. The GPU can do those flops in microseconds and spends milliseconds fetching.

Attention is the same story with worse constants. The score matrix is N×N. At a sequence length of 8192, one head materializes 64 million values — write them out, read them back for softmax, write again, read again for the second matmul. FlashAttention's entire contribution is refusing to do that.

Level 1 · what you configure stays in registers

Fuse the epilogue

When the GEMM finishes a tile, the results are sitting in accumulator registers. Anything elementwise you apply right there is free — the data never moves. cuBLASLt exposes this as a descriptor attribute, so it costs you three lines.

cuBLASLt · GEMM + bias + ReLU, one kernelC++

cublasLtMatmulDescCreate(&op, CUBLAS_COMPUTE_32F, CUDA_R_32F);

// the three lines that matter
cublasLtEpilogue_t epi = CUBLASLT_EPILOGUE_RELU_BIAS;
cublasLtMatmulDescSetAttribute(op, CUBLASLT_MATMUL_DESC_EPILOGUE,
                               &epi, sizeof(epi));
cublasLtMatmulDescSetAttribute(op, CUBLASLT_MATMUL_DESC_BIAS_POINTER,
                               &dBias, sizeof(dBias));

// layouts, heuristic, then run — bias and ReLU ride along for free
cublasLtMatmulAlgoGetHeuristic(lt, op, Lb, La, Ld, Ld, pref, 1, &heur, &found);
cublasLtMatmul(lt, op, &alpha, dB, Lb, dA, La, &beta, dD, Ld, dD, Ld,
               &heur.algo, ws, wsSize, stream);

Epilogue	Fuses
_BIAS	Per-column bias add
_RELU_BIAS	Bias then ReLU
_GELU_BIAS	Bias then GELU
_GELU_AUX_BIAS	Same, plus saves pre-activation for the backward pass
_DGELU, _DRELU	Backward: activation gradient fused into the GEMM
_BGRADA, _BGRADB	Bias gradient reduction fused into the GEMM

CUTLASS gives you the same thing as a template parameter, and in 3.x an Epilogue Visitor Tree that composes arbitrary elementwise DAGs:

CUTLASS 2.x · GELU in the epilogueC++

using Epilogue = cutlass::epilogue::thread::LinearCombinationGELU<
    float,                                     // output element
    128 / cutlass::sizeof_bits<float>::value,  // vector width
    float, float>;                             // accum, compute

using Gemm = cutlass::gemm::device::Gemm<
    cutlass::half_t, cutlass::layout::RowMajor,
    cutlass::half_t, cutlass::layout::RowMajor,
    float,           cutlass::layout::RowMajor,
    float, cutlass::arch::OpClassTensorOp, cutlass::arch::Sm80,
    cutlass::gemm::GemmShape<128,128,32>,   // threadblock tile
    cutlass::gemm::GemmShape<64,64,32>,     // warp tile
    cutlass::gemm::GemmShape<16,8,16>,      // mma instruction
    Epilogue>;

Level 2 · what you write stays in registers

Your own epilogue, in the accumulators

If the activation you need isn't on the list, reach into the WMMA fragment directly. After the K loop the accumulators still hold your tile — apply the function elementwise before storing.

Safe: uniform elementwise on a fragmentCUDA C++

__device__ inline float gelu(float x) {
    const float k = 0.7978845608f;          // sqrt(2/pi)
    return 0.5f * x * (1.0f + tanhf(k * (x + 0.044715f * x*x*x)));
}

// ... K loop has ended, acc[][] is live in registers ...
for (int i = 0; i < WM/16; ++i)
for (int j = 0; j < WN/16; ++j) {
    // every element gets identical treatment, so the opaque
    // fragment layout is irrelevant. this is the one guarantee.
    #pragma unroll
    for (int e = 0; e < acc[i][j].num_elements; ++e)
        acc[i][j].x[e] = gelu(acc[i][j].x[e]);

    wmma::store_matrix_sync(&C[row*N + col], acc[i][j], N,
                            wmma::mem_row_major);
}

The trap

A bias is not uniform — it varies by column, and you cannot know which column acc.x[e] holds. The fragment layout is undocumented and differs across architectures. Code that guesses it works on Ampere and silently corrupts on Hopper.

Portable fix for position-dependent opsCUDA C++

// stage through shared memory, where indices are known.
// +1 skew keeps the 32 lanes off the same bank.
__shared__ float smem[WARPS][16][17];
float (*tile)[17] = smem[warpId];

wmma::store_matrix_sync(&tile[0][0], acc[i][j], 17, wmma::mem_row_major);
__syncwarp();

for (int e = lane; e < 16*16; e += 32) {
    int r = e / 16, c = e % 16;
    tile[r][c] = gelu(tile[r][c] + bias[col + c]);   // index is real now
}
__syncwarp();
// then write tile -> C with coalesced 128-bit stores

You paid one SRAM round trip instead of two HBM round trips. Still a 10× saving on the traffic that mattered.

Level 3 · what you restructure stays in SRAM

Fuse the loop: online softmax

Epilogue fusion only handles elementwise work. Softmax needs a row max and a row sum — a reduction across the whole key dimension — which looks like it forces you to materialize the full row before you can normalize it.

It doesn't. Keep a running max and a running sum, and rescale the partial output whenever the max moves:

for each block j of keys: S_j = Q · K_jᵀ · scale m_new = max(m_old, rowmax(S_j)) α = exp(m_old − m_new) ← correction factor P_j = exp(S_j − m_new) ℓ_new = α·ℓ_old + rowsum(P_j) O_new = α·O_old + P_j · V_j once, at the end: O = O / ℓ

Each block is exact given what it has seen; α retroactively fixes the earlier blocks. The N×N scores exist only as a tile in SRAM. FlashAttention-2's refinement is visible above — the division happens once in the epilogue, not per block, because non-matmul flops are the expensive kind here.

The real constraint

On A100, tensor cores do ~312 TFLOP/s while the units running exp and the row reductions do ~19. A few percent of non-matmul flops can eat half your runtime. That ratio, not memory, is what shapes every version after FA1.

In CUDA this is a few thousand lines. In Triton it is legible:

Fused attention forward · single head, non-causalTriton

@triton.jit
def flash_attn_fwd(Q, K, V, O, sm_scale,
                   N_CTX: tl.constexpr, D: tl.constexpr,
                   BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr):
    start_m = tl.program_id(0)
    offs_m  = start_m * BLOCK_M + tl.arange(0, BLOCK_M)
    offs_d  = tl.arange(0, D)

    # this Q block is loaded once and never reread
    q = tl.load(Q + offs_m[:, None] * D + offs_d[None, :])

    m_i = tl.full([BLOCK_M], float("-inf"), tl.float32)
    l_i = tl.zeros([BLOCK_M], tl.float32)
    acc = tl.zeros([BLOCK_M, D], tl.float32)

    for start_n in range(0, N_CTX, BLOCK_N):
        offs_n = start_n + tl.arange(0, BLOCK_N)
        k = tl.load(K + offs_n[:, None] * D + offs_d[None, :])
        v = tl.load(V + offs_n[:, None] * D + offs_d[None, :])

        # ---- MMA 1: tensor cores ----
        qk = tl.dot(q, tl.trans(k)) * sm_scale

        # ---- online softmax: CUDA cores + SFU, the slow lane ----
        m_new = tl.maximum(m_i, tl.max(qk, axis=1))
        alpha = tl.exp(m_i - m_new)
        p     = tl.exp(qk - m_new[:, None])

        acc   = acc * alpha[:, None]        # retro-fix earlier blocks
        l_i   = l_i * alpha + tl.sum(p, axis=1)
        m_i   = m_new

        # ---- MMA 2: tensor cores ----
        acc  += tl.dot(p.to(v.dtype), v)

    acc = acc / l_i[:, None]                # one division, in the epilogue
    tl.store(O + offs_m[:, None] * D + offs_d[None, :], acc)

tl.dot lowers straight to tensor-core MMA. Triton picks the shared-memory layouts, swizzling, and async pipelining for you — which is most of the reason this fits on a screen.

FlashAttention-3 attacks the slow lane directly. On Hopper it warp-specializes: some warps do nothing but wgmma, others do nothing but softmax, and a "pingpong" schedule runs block j's softmax concurrently with block j+1's GEMM. Roughly 740 TFLOP/s in FP16, ~1.2 PFLOP/s in FP8.

Level 4 · what you delegate

Let a compiler do it

torch.compile — fuses pointwise chains automatically via Inductor, which emits Triton. Zero effort, and it catches the boring 80%: bias, activations, dropout, residual adds, layernorm tails.
Triton — when you need a fused op Inductor won't find, or a custom attention variant. You write blocks; it writes the memory plumbing.
CUTLASS / CuTe — when you need control over the pipeline itself: unusual tiling, mixed dtypes per leg, a fusion that crosses the mainloop rather than the epilogue.
Raw CUDA + PTX — when you are chasing the last 15% and know exactly which instruction is stalling.

Work down that list, not up it. Most "I need a custom kernel" turns out to be an epilogue.

Limits

Where fusion stops paying

Register pressure. Fusing more means holding more live values. Past the limit the compiler spills to "local" memory — which is HBM wearing a disguise. Fusion that spills is slower than no fusion. Watch local_load / local_store in Nsight Compute.
Occupancy collapse. Bigger tiles and fatter epilogues cut the number of resident blocks. Fewer blocks means less latency hiding. There is a real optimum and it is rarely "fuse everything."
Reductions that aren't online. Softmax and LayerNorm have online formulations (running max/sum, Welford), so they fuse. Split-K GEMM genuinely needs a device-wide sync — a second kernel or atomics.
Backward passes. Fusing forward means not storing intermediates, so backward has to recompute them. FlashAttention recomputes the scores. That is a real flop cost traded against real memory savings, and which side wins depends on sequence length.

Verification

Confirm it actually fused

Nsight Computeshell

ncu --metrics \
  sm__inst_executed_pipe_tensor.sum,\
  dram__bytes.sum,\
  l1tex__t_bytes.sum,\
  launch__registers_per_thread \
  ./your_binary

sm__inst_executed_pipe_tensor.sum is zero → you never touched a tensor core. Check alignment and compute type before anything else.
dram__bytes.sum should drop by roughly the number of intermediates you eliminated. If it didn't, the fusion didn't happen.
launch__registers_per_thread near 255 → you are about to spill. Back off.