Epilogue Subtiling Example

This example demonstrates matmul kernels with heavy epilogues that benefit from epilogue subtiling on Blackwell (sm_100+). Epilogue subtiling splits the store from [BLOCK_M, BLOCK_N] into SUBTILE_FACTOR x [BLOCK_M, BLOCK_N / SUBTILE_FACTOR], halving the accumulator shared-memory footprint and enabling an extra pipeline stage.

Imports

from __future__ import annotations

import torch

import helion
from helion._testing import DEVICE
from helion._testing import HALF_DTYPE
from helion._testing import run_example
import helion.language as hl

Kernel 1 – Matmul + Residual + Bias + GELU + Cast

CUTLASS-style residual + bias + GELU forward epilogue with two fp32 reads (residual, bias) fused into the output tile.

@helion.kernel(static_shapes=True)
def matmul_bias_residual_gelu_cast(
    x: torch.Tensor,
    w: torch.Tensor,
    bias: torch.Tensor,
    residual: torch.Tensor,
) -> torch.Tensor:
    m, k = x.size()
    _, n = w.size()
    out = torch.empty([m, n], dtype=HALF_DTYPE, device=x.device)

    for tile_m, tile_n in hl.tile([m, n]):
        acc = hl.zeros([tile_m, tile_n], dtype=torch.float32)
        for tile_k in hl.tile(k):
            acc = torch.addmm(acc, x[tile_m, tile_k], w[tile_k, tile_n])

        val = acc * 1.25
        val = val + residual[tile_m, tile_n].to(torch.float32) * 0.5
        val = val + bias[tile_n]
        val = torch.nn.functional.gelu(val)
        out[tile_m, tile_n] = val.to(HALF_DTYPE)

    return out

Kernel 2 – Matmul + Bias + GELU with Auxiliary Output

cuBLASLt / CUTLASS-style GELU+AUX forward epilogue that writes both the pre-activation (aux) and post-GELU (out) tensors.

@helion.kernel(static_shapes=True)
def matmul_bias_gelu_aux(
    x: torch.Tensor,
    w: torch.Tensor,
    bias: torch.Tensor,
) -> tuple[torch.Tensor, torch.Tensor]:
    m, k = x.size()
    _, n = w.size()
    out = torch.empty([m, n], dtype=HALF_DTYPE, device=x.device)
    aux = torch.empty([m, n], dtype=HALF_DTYPE, device=x.device)

    for tile_m, tile_n in hl.tile([m, n]):
        acc = hl.zeros([tile_m, tile_n], dtype=torch.float32)
        for tile_k in hl.tile(k):
            acc = torch.addmm(acc, x[tile_m, tile_k], w[tile_k, tile_n])

        pre = acc * 1.25
        pre = pre + bias[tile_n]
        aux[tile_m, tile_n] = pre.to(HALF_DTYPE)
        out[tile_m, tile_n] = torch.nn.functional.gelu(pre).to(HALF_DTYPE)

    return out, aux

Verification

def check(m: int, k: int, n: int) -> None:
    x = torch.randn([m, k], device=DEVICE, dtype=HALF_DTYPE)
    w = torch.randn([k, n], device=DEVICE, dtype=HALF_DTYPE)
    bias = torch.randn([n], device=DEVICE, dtype=HALF_DTYPE)
    residual = torch.randn([m, n], device=DEVICE, dtype=HALF_DTYPE)

    def baseline_residual_gelu(
        x: torch.Tensor,
        w: torch.Tensor,
        bias: torch.Tensor,
        residual: torch.Tensor,
    ) -> torch.Tensor:
        acc = x.float() @ w.float()
        val = acc * 1.25 + residual.float() * 0.5 + bias.float()
        return torch.nn.functional.gelu(val).to(HALF_DTYPE)

    run_example(
        matmul_bias_residual_gelu_cast,
        baseline_residual_gelu,
        (x, w, bias, residual),
    )

    def baseline_gelu_aux(
        x: torch.Tensor,
        w: torch.Tensor,
        bias: torch.Tensor,
    ) -> tuple[torch.Tensor, torch.Tensor]:
        acc = x.float() @ w.float()
        pre = acc * 1.25 + bias.float()
        return torch.nn.functional.gelu(pre).to(HALF_DTYPE), pre.to(HALF_DTYPE)

    run_example(
        matmul_bias_gelu_aux,
        baseline_gelu_aux,  # pyrefly: ignore[bad-argument-type]
        (x, w, bias),
    )

Main

def main() -> None:
    check(8192, 8192, 8192)


if __name__ == "__main__":
    main()