Route matmul_kernel_tma

src/fused_mm_sampling/persistent_matmul.py:None–None

triton.autotune(
    configs=matmul_get_configs(pre_hook=matmul_tma_set_block_size_hook),
    key=["M", "N", "K", "WARP_SPECIALIZE"],
)

import argparse
import itertools
from contextlib import contextmanager
from typing import Optional