Function torch_matmul

src/fused_mm_sampling/persistent_matmul.py:686–696

(a, b)

def torch_matmul(a, b):
    M, K = a.shape
    N, K = b.shape
    bytes_per_elem = a.element_size()
    flops_str = f"flops{bytes_per_elem * 8}"
    with proton.scope(
        f"torch [M={M}, N={N}, K={K}]",
        {"bytes": bytes_per_elem * (M * K + N * K + M * N), flops_str: 2.0 * M * N * K},
    ):
        c = torch.matmul(a, b.T)
    return c


def torch_matmul_nontransposed(a, b):