File size: 797 Bytes
d91766b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 | import itertools
def build_chunked_prefill_configs():
"""Autotune configs for chunked prefill unified kernel.
Note: The autotune key is
["NUM_GROUPS", "HEAD_DIM", "PAGE_SIZE", "DLLM_BLOCK_SIZE", "IS_BLOCK_CAUSAL", "IS_PREFIX_FULL"],
so configs are selected based on those runtime parameters.
"""
BLOCK_M_LIST = [64, 128]
BLOCK_N_LIST = [64, 128]
NUM_STAGES_LIST = [1, 2, 3]
NUM_WARPS_LIST = [4, 8]
CONFIGS = list(
itertools.product(
BLOCK_M_LIST,
BLOCK_N_LIST,
NUM_STAGES_LIST,
NUM_WARPS_LIST,
)
)
return [
{
"BLOCK_M": c[0],
"BLOCK_N": c[1],
"num_stages": c[2],
"num_warps": c[3],
}
for c in CONFIGS
]
|