@triton.autotune(configs=[triton.Config({'BLOCK_SIZE_B': 32, 'BLOCK_SIZE_E':
    32}, num_stages=3, num_warps=4), triton.Config({'BLOCK_SIZE_B': 64,
    'BLOCK_SIZE_E': 32}, num_stages=2, num_warps=4), triton.Config({
    'BLOCK_SIZE_B': 32, 'BLOCK_SIZE_E': 64}, num_stages=2, num_warps=4),
    triton.Config({'BLOCK_SIZE_B': 64, 'BLOCK_SIZE_E': 64}, num_stages=2,
    num_warps=4)], key=['H', 'B', 'D', 'E'])
@triton.jit
def mlp_wide_kernel_bwd2(x_ptr, w1_ptr, w2_ptr, o_ptr, dx_ptr, dw1_ptr,
    dw2_ptr, do_ptr, H, B, D: tl.constexpr, E, stride_xb, stride_xd,
    stride_w1d, stride_w1e, stride_w2e, stride_w2d, stride_ob, stride_od,
    stride_dxb, stride_dxd, stride_dw1d, stride_dw1e, stride_dw2e,
    stride_dw2d, stride_dob, stride_dod, 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)
    - dX has shape (B, D)
    - dW1 has shape (D, E)
    - dW2 has shape (E, D)
    - dO has shape (B, D)
    """
    pid = tl.program_id(axis=0)
    pid_x_w = 0
    batch_groups_e = tl.cdiv(E, BLOCK_SIZE_E)
    batch_groups_b = tl.cdiv(B, BLOCK_SIZE_B)
    idx = pid % (batch_groups_e + batch_groups_b)
    pid_h = pid // (batch_groups_e + batch_groups_b)
    TARGET_TYPE = x_ptr.type.element_ty
    offs_b = tl.arange(0, BLOCK_SIZE_B)
    offs_d = tl.arange(0, D)
    offs_e = tl.arange(0, BLOCK_SIZE_E)
    if idx >= batch_groups_e:
        pid_b = idx - batch_groups_e
        dx_ptrs = dx_ptr + ((pid_h * B + pid_b * BLOCK_SIZE_B + offs_b[:,
            None]) * stride_dxb + offs_d[None, :] * stride_dxd)
        dx_mask = (offs_b[:, None] < B - pid_b * BLOCK_SIZE_B) & (offs_d[
            None, :] < D)
        dx = tl.zeros((BLOCK_SIZE_B, D), dtype=tl.float32)
        dx = _mlp_wide_kernel_bwd_dx(dx, pid_h, pid_b, x_ptr, w1_ptr,
            w2_ptr, o_ptr, dx_ptr, dw1_ptr, dw2_ptr, do_ptr, H, B, D, E,
            stride_xb, stride_xd, stride_w1d, stride_w1e, stride_w2e,
            stride_w2d, stride_ob, stride_od, stride_dxb, stride_dxd,
            stride_dw1d, stride_dw1e, stride_dw2e, stride_dw2d, stride_dob,
            stride_dod, BLOCK_SIZE_B, BLOCK_SIZE_E, ACTIVATION)
        tl.store(dx_ptrs, dx.to(TARGET_TYPE), mask=dx_mask)
    else:
        pid_e = idx
        dw1_ptrs = dw1_ptr + ((pid_h * D + offs_d[:, None]) * stride_dw1d +
            (pid_e * BLOCK_SIZE_E + offs_e[None, :]) * stride_dw1e)
        dw1_mask = (offs_d[:, None] < D) & (offs_e[None, :] < E - pid_e *
            BLOCK_SIZE_E)
        dw2_ptrs = dw2_ptr + ((pid_h * E + pid_e * BLOCK_SIZE_E + offs_e[:,
            None]) * stride_dw2e + offs_d[None, :] * stride_dw2d)
        dw2_mask = (offs_e[:, None] < E - pid_e * BLOCK_SIZE_E) & (offs_d[
            None, :] < D)
        dw1 = tl.zeros((D, BLOCK_SIZE_E), dtype=tl.float32)
        dw2 = tl.zeros((BLOCK_SIZE_E, D), dtype=tl.float32)
        dw1, dw2 = _mlp_wide_kernel_bwd_dw1w2(dw1, dw2, pid_h, pid_e, x_ptr,
            w1_ptr, w2_ptr, o_ptr, dx_ptr, dw1_ptr, dw2_ptr, do_ptr, H, B,
            D, E, stride_xb, stride_xd, stride_w1d, stride_w1e, stride_w2e,
            stride_w2d, stride_ob, stride_od, stride_dxb, stride_dxd,
            stride_dw1d, stride_dw1e, stride_dw2e, stride_dw2d, stride_dob,
            stride_dod, BLOCK_SIZE_B, BLOCK_SIZE_E, ACTIVATION)
        tl.store(dw1_ptrs, dw1.to(TARGET_TYPE), mask=dw1_mask)
        tl.store(dw2_ptrs, dw2.to(TARGET_TYPE), mask=dw2_mask)
