benchmark_configs.py 82 Bytes
Newer Older
1
2
# BATCH, N_HEADS, SEQ_LEN, D_HEAD, TOPK, BLOCK
configs = [[4, 2, 256, 64, 2, 64]]