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:
Claude
2026-08-03 11:18:18 +00:00
parent fe64650681
commit a2a5dd8f00
3 changed files with 108 additions and 87 deletions

View File

@@ -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'],
)