@triton.autotune(configs=[triton.Config({'BLOCK_SIZE_B': 32, 'BLOCK_SIZE_E':
    32}, num_stages=4, num_warps=4), triton.Config({'BLOCK_SIZE_B': 64,
    'BLOCK_SIZE_E': 32}, num_stages=4, num_warps=4), triton.Config({
    'BLOCK_SIZE_B': 32, 'BLOCK_SIZE_E': 64}, num_stages=3, num_warps=4),
    triton.Config({'BLOCK_SIZE_B': 64, 'BLOCK_SIZE_E': 64}, num_stages=3,
    num_warps=4)], key=['B', 'D', 'E'])
@triton.jit
def mlp_wide_kernel_fwd(x_ptr, w1_ptr, w2_ptr, o_ptr, H, B, D: tl.constexpr,
    E, stride_xb, stride_xd, stride_w1d, stride_w1e, stride_w2e, stride_w2d,
    stride_ob, stride_od, BLOCK_SIZE_B: tl.constexpr, BLOCK_SIZE_E: tl.
    constexpr, ACTIVATION: tl.constexpr):
    """Kernel for computing the mlp
    Z = X @ W1, H = f(Z), O = H @ W2.
    - X has shape (B, D)
    - W1 has shape (D, E)
    - W2 has shape (E, D)
    - O has shape (B, D)
    """
    pid = tl.program_id(axis=0)
    batch_groups = tl.cdiv(B, BLOCK_SIZE_B)
    pid_b = pid % batch_groups
    pid_h = pid // batch_groups
    TARGET_TYPE = x_ptr.type.element_ty
    x_ptrs = tl.make_block_ptr(base=x_ptr, shape=(B * H, D), strides=(
        stride_xb, stride_xd), offsets=(pid_h * B + pid_b * BLOCK_SIZE_B, 0
        ), block_shape=(BLOCK_SIZE_B, D), order=(1, 0))
    w1_ptrs = tl.make_block_ptr(base=w1_ptr, shape=(D * H, E), strides=(
        stride_w1d, stride_w1e), offsets=(pid_h * D, 0), block_shape=(D,
        BLOCK_SIZE_E), order=(1, 0))
    w2_ptrs = tl.make_block_ptr(base=w2_ptr, shape=(E * H, D), strides=(
        stride_w2e, stride_w2d), offsets=(pid_h * E, 0), block_shape=(
        BLOCK_SIZE_E, D), order=(1, 0))
    o_ptrs = tl.make_block_ptr(base=o_ptr, shape=(B * H, D), strides=(
        stride_ob, stride_od), offsets=(pid_h * B + pid_b * BLOCK_SIZE_B, 0
        ), block_shape=(BLOCK_SIZE_B, D), order=(1, 0))
    x = tl.load(x_ptrs)
    o = tl.zeros((BLOCK_SIZE_B, D), dtype=tl.float32)
    for e in range(0, tl.cdiv(E, BLOCK_SIZE_E)):
        w1 = tl.load(w1_ptrs)
        w2 = tl.load(w2_ptrs)
        z = tl.dot(x, w1, out_dtype=tl.float32)
        if ACTIVATION == 'leaky_relu':
            z = leaky_relu(z).to(TARGET_TYPE)
        elif ACTIVATION == 'silu':
            z = silu(z).to(TARGET_TYPE)
        elif ACTIVATION == 'sigmoid':
            z = tl.sigmoid(z).to(TARGET_TYPE)
        else:
            z = z.to(TARGET_TYPE)
        o = tl.dot(z, w2, o, out_dtype=tl.float32)
        w1_ptrs = tl.advance(w1_ptrs, (0, BLOCK_SIZE_E))
        w2_ptrs = tl.advance(w2_ptrs, (BLOCK_SIZE_E, 0))
    o = o.to(TARGET_TYPE)
    tl.store(o_ptrs, o)
