diff --git a/vllm/attention/ops/triton_flash_attention.py b/vllm/attention/ops/triton_flash_attention.py index aa7414ac..8ada332c 100644 --- a/vllm/attention/ops/triton_flash_attention.py +++ b/vllm/attention/ops/triton_flash_attention.py @@ -378,6 +378,48 @@ def _attn_fwd_inner( num_stages=2, num_warps=2, ), + # BI-V100 warp-uniform optimization (from make_warp_uniform.cuh): + # With 16 SMs and small CTAs (num_warps=2 → 64 threads), we can + # run 16+ CTAs simultaneously. All threads in each warp access + # the same batch/head (uniform control flow) → zero divergence. + # num_stages=2 doubles prefetch window → matches 64KB BIF sweet spot. + # PRE_LOAD_V=True: pre-load V tile into registers before score + # computation. Safe for small BLOCK_N because register pressure is: + # Q: BLOCK_M×BLOCK_D = 32×256 = 8K regs (fp16) + # V: BLOCK_N×BLOCK_D = 32×256 = 8K regs (fp16) + # Total: 16K regs << 64K regs/SM available on BI-V100 + # This mirrors CCCL agent_reduce ConsumeFullTile vectorized path + # which loads VectorT into registers before applying reduction. + triton.Config( + { + "BLOCK_M": 32, + "BLOCK_N": 32, + "waves_per_eu": 4, + "PRE_LOAD_V": True, + }, + num_stages=2, + num_warps=2, + ), + triton.Config( + { + "BLOCK_M": 64, + "BLOCK_N": 32, + "waves_per_eu": 2, + "PRE_LOAD_V": True, + }, + num_stages=2, + num_warps=4, + ), + triton.Config( + { + "BLOCK_M": 16, + "BLOCK_N": 32, + "waves_per_eu": 4, + "PRE_LOAD_V": True, + }, + num_stages=2, + num_warps=2, + ), ], key=['IS_CAUSAL', 'dropout_p', 'BLOCK_DMODEL'], )