feat: asymmetric BLOCK_M/BLOCK_N search + re-add BI-V100 autotune configs
bench_triton_prefill.py:
- Split --block into --block (BLOCK_M) and --block-n (BLOCK_N)
- Each (M, N, warps) combo triggers Triton JIT recompilation
- Enables finding asymmetric optima like M=64,N=32 that save SMEM
triton_flash_attention.py:
- Re-add 3 BI-V100 autotune configs (64x32, 32x64, 64x64 with warps=4)
- These were wrongly reverted in 8c1955d -- autotune is zero-risk
run_on_bi100.sh:
- Updated to use asymmetric block search
This commit is contained in:
@@ -302,6 +302,39 @@ def _attn_fwd_inner(
|
||||
num_stages=1,
|
||||
num_warps=4,
|
||||
),
|
||||
# BI-V100 (SM=16, SMEM≤48KB): smaller BLOCK_M maintains occupancy,
|
||||
# asymmetric M/N trades Q-tile for longer K/V sweeps per CTA.
|
||||
# Autotune will discard these if they're slower — zero risk.
|
||||
triton.Config(
|
||||
{
|
||||
"BLOCK_M": 64,
|
||||
"BLOCK_N": 32,
|
||||
"waves_per_eu": 2,
|
||||
"PRE_LOAD_V": False,
|
||||
},
|
||||
num_stages=1,
|
||||
num_warps=4,
|
||||
),
|
||||
triton.Config(
|
||||
{
|
||||
"BLOCK_M": 32,
|
||||
"BLOCK_N": 64,
|
||||
"waves_per_eu": 2,
|
||||
"PRE_LOAD_V": False,
|
||||
},
|
||||
num_stages=1,
|
||||
num_warps=4,
|
||||
),
|
||||
triton.Config(
|
||||
{
|
||||
"BLOCK_M": 64,
|
||||
"BLOCK_N": 64,
|
||||
"waves_per_eu": 2,
|
||||
"PRE_LOAD_V": False,
|
||||
},
|
||||
num_stages=1,
|
||||
num_warps=4,
|
||||
),
|
||||
],
|
||||
key=['IS_CAUSAL', 'dropout_p', 'BLOCK_DMODEL'],
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user