@triton.autotune(configs=[triton.Config({'BLOCK_SIZE_M': 256,
    'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 4, 'grf_mode':
    'large'}, num_stages=2, num_warps=32), triton.Config({'BLOCK_SIZE_M': 
    256, 'BLOCK_SIZE_N': 256, 'BLOCK_SIZE_K': 32, 'GROUP_SIZE_M': 4,
    'grf_mode': 'large'}, num_stages=3, num_warps=32), triton.Config({
    'BLOCK_SIZE_M': 256, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K': 32,
    'GROUP_SIZE_M': 4, 'grf_mode': 'large'}, num_stages=2, num_warps=32),
    triton.Config({'BLOCK_SIZE_M': 64, 'BLOCK_SIZE_N': 128, 'BLOCK_SIZE_K':
    32, 'GROUP_SIZE_M': 4, 'grf_mode': 'large'}, num_stages=2, num_warps=32
    ), triton.Config({'BLOCK_SIZE_M': 8, 'BLOCK_SIZE_N': 512,
    'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 1, 'grf_mode': 'large'}, num_stages
    =2, num_warps=32), triton.Config({'BLOCK_SIZE_M': 8, 'BLOCK_SIZE_N': 
    128, 'BLOCK_SIZE_K': 64, 'GROUP_SIZE_M': 1, 'grf_mode': 'large'},
    num_stages=2, num_warps=4)], key=['M', 'N', 'K'])
@triton.jit
def matmul_kernel_with_block_pointers_batched(a_ptr, b_ptr, c_ptr, B: tl.
    constexpr, M: tl.constexpr, N: tl.constexpr, K: tl.constexpr, stride_az:
    tl.constexpr, stride_am: tl.constexpr, stride_ak: tl.constexpr,
    stride_bz: tl.constexpr, stride_bk: tl.constexpr, stride_bn: tl.
    constexpr, stride_cz: tl.constexpr, stride_cm: tl.constexpr, stride_cn:
    tl.constexpr, BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr,
    BLOCK_SIZE_K: tl.constexpr, GROUP_SIZE_M: tl.constexpr):
    bid = tl.program_id(axis=0)
    pid = tl.program_id(axis=1)
    num_pid_m = tl.cdiv(M, BLOCK_SIZE_M)
    num_pid_n = tl.cdiv(N, BLOCK_SIZE_N)
    num_pid_in_group = GROUP_SIZE_M * num_pid_n
    group_id = pid // num_pid_in_group
    first_pid_m = group_id * GROUP_SIZE_M
    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)
    pid_m = first_pid_m + pid % group_size_m
    pid_n = pid % num_pid_in_group // group_size_m
    offset_a = bid.to(tl.int64) * stride_az
    offset_b = bid.to(tl.int64) * stride_bz
    a_block_ptr = tl.make_block_ptr(base=a_ptr + offset_a, shape=(M, K),
        strides=(stride_am, stride_ak), offsets=(pid_m * BLOCK_SIZE_M, 0),
        block_shape=(BLOCK_SIZE_M, BLOCK_SIZE_K), order=(1, 0))
    b_block_ptr = tl.make_block_ptr(base=b_ptr + offset_b, shape=(K, N),
        strides=(stride_bk, stride_bn), offsets=(0, pid_n * BLOCK_SIZE_N),
        block_shape=(BLOCK_SIZE_K, BLOCK_SIZE_N), order=(1, 0))
    accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)
    for _ in range(0, K, BLOCK_SIZE_K):
        a = tl.load(a_block_ptr, boundary_check=(0, 1))
        a = a.to(tl.float32)
        a = tl.math.exp(a)
        a = a.to(tl.bfloat16)
        b = tl.load(b_block_ptr, boundary_check=(0, 1))
        accumulator += tl.dot(a, b)
        a_block_ptr = tl.advance(a_block_ptr, (0, BLOCK_SIZE_K))
        b_block_ptr = tl.advance(b_block_ptr, (BLOCK_SIZE_K, 0))
    c = accumulator.to(tl.float32)
    offset_c = bid.to(tl.int64) * stride_cz
    c_block_ptr = tl.make_block_ptr(base=c_ptr + offset_c, shape=(M, N),
        strides=(stride_cm, stride_cn), offsets=(pid_m * BLOCK_SIZE_M, 
        pid_n * BLOCK_SIZE_N), block_shape=(BLOCK_SIZE_M, BLOCK_SIZE_N),
        order=(1, 0))
    tl.store(c_block_ptr, c, boundary_check=(0, 1))
