@triton.jit
def _matmul_kernel(A, B, C, M, N, K, **meta):
    TILE_M = meta['BLOCK_M']
    TILE_N = meta['BLOCK_N']
    TILE_K = 128
    m = tl.program_id(0) * TILE_M + tl.arange(0, TILE_M)
    n = tl.program_id(1) * TILE_N + tl.arange(0, TILE_N)
    acc = tl.zeros((TILE_M, TILE_N), dtype=tl.float32)
    for k in range(0, K, TILE_K):
        a = tl.load(A + m[:, None] * K + k, mask=[m[:, None] < M, None],
            other=0.0)
        b = tl.load(B + k * N + n, mask=[None, n < N], other=0.0)
        acc += tl.dot(a, b)
    tl.store(C + m[:, None] * N + n, acc, mask=[m[:, None] < M, n < N])
