@triton.autotune(configs=[triton.Config({'BLOCK_ROW': 16, 'BLOCK_COL': 16},
    num_stages=5, num_warps=1), triton.Config({'BLOCK_ROW': 32, 'BLOCK_COL':
    32}, num_stages=5, num_warps=1), triton.Config({'BLOCK_ROW': 64,
    'BLOCK_COL': 32}, num_stages=5, num_warps=2), triton.Config({
    'BLOCK_ROW': 32, 'BLOCK_COL': 64}, num_stages=5, num_warps=2), triton.
    Config({'BLOCK_ROW': 128, 'BLOCK_COL': 64}, num_stages=4, num_warps=4),
    triton.Config({'BLOCK_ROW': 64, 'BLOCK_COL': 128}, num_stages=4,
    num_warps=4), triton.Config({'BLOCK_ROW': 128, 'BLOCK_COL': 128},
    num_stages=4, num_warps=4)], key=['M', 'N', 'K'])
@triton.jit
def kernel_fma(OUT, ACT_INPUTS, INPUT, WEIGHT, BIAS, M, N, K, stride_om,
    stride_im, stride_wn, stride_wk, **META):
    """
    Kernel for computing Out = activation(A x W + C)

    - Input has shape (M, K)
    - Weight has shape (K, N)
    - Bias has shape (N,)
    - Output has shape (M, N)
    - ActInputs (optional) has shape (M, N)

    'ActInputs' optionally saves the A x W + C intermediate for backward computations

    This kernel will consolidate over K
    """
    BLOCK_M, GROUP_M = META['BLOCK_ROW'], META['GROUP_ROW']
    BLOCK_N, BLOCK_K = META['BLOCK_COL'], META['BLOCK_K']
    pid = tl.program_id(axis=0)
    num_pid_m = tl.cdiv(M, BLOCK_M)
    num_pid_n = tl.cdiv(N, BLOCK_N)
    num_pid_in_group = GROUP_M * num_pid_n
    group_id = pid // num_pid_in_group
    first_pid_m = group_id * GROUP_M
    GROUP_M = min(num_pid_m - first_pid_m, GROUP_M)
    pid_m = first_pid_m + pid % GROUP_M
    pid_n = pid % num_pid_in_group // GROUP_M
    rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
    rn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
    rk = tl.arange(0, BLOCK_K)
    input_ptrs = INPUT + rm[:, None] * stride_im + rk[None, :]
    weight_ptrs = WEIGHT + rk[:, None] * stride_wk + rn[None, :] * stride_wn
    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
    if META['BIAS']:
        bias = tl.load(BIAS + rn, mask=rn < N, other=0.0).to(tl.float32)
        acc += bias[None, :]
    for _ in range(K, 0, -BLOCK_K):
        a = tl.load(input_ptrs, mask=(rk[None, :] < K) & (rm[:, None] < M),
            other=0.0)
        w = tl.load(weight_ptrs, mask=(rk[:, None] < K) & (rn[None, :] < N),
            other=0.0)
        acc += tl.dot(a, w).to(tl.float32)
        input_ptrs += BLOCK_K
        weight_ptrs += BLOCK_K * stride_wk
    if META['SAVE_ACT_INPUTS']:
        act_in_ptrs = ACT_INPUTS + rm[:, None] * stride_om + rn[None, :]
        tl.store(act_in_ptrs, acc, mask=(rm[:, None] < M) & (rn[None, :] < N))
    if META['ACTIVATION']:
        acc = META['ACTIVATION'](acc)
    out_ptrs = OUT + rm[:, None] * stride_om + rn[None, :]
    tl.store(out_ptrs, acc, mask=(rm[:, None] < M) & (rn[None, :] < N))
