[flash_attn] CCCL-derived autotune configs: num_stages=2 + small-tile
Two findings from CCCL benchmarks applied to Triton autotune configs:
1. num_stages=2 (from transform bif=8 finding):
CCCL transform benchmark (babelstream.cu) search space includes
TUNE_BIF_BIAS from -16 to +16. BI-V100 bench found bif=8 (64KB
prefetch window) dominates across all problem sizes. Physical basis:
BW_per_SM × memory_latency = 56 GB/s × 1100ns ≈ 62KB
Triton's num_stages is the software pipelining equivalent of CCCL's
bytes_in_flight. num_stages=2 doubles the prefetch window from ~32KB
to ~64KB, matching the optimal BW×latency product.
2. Small-tile high-occupancy (from scan no_delay finding):
CCCL scan benchmark (sum.cu) found dcid=0 (no_delay) optimal on
BI-V100 because 16 SMs produce only ~32 CTAs, so the tile_status
array fits entirely in 6MB L2 with zero inter-CTA contention.
Implication: more smaller CTAs can saturate the 16 SMs better than
fewer large CTAs, especially for short sequences.
Added 3 new configs, all with num_stages=2 or waves_per_eu=4.
Triton autotune will select the fastest; no risk of regression.
Source: cccl_upstream/cub/benchmarks/bench/transform/babelstream.cu
cccl_upstream/cub/benchmarks/bench/scan/exclusive/sum.cu
This commit is contained in:
@@ -335,6 +335,49 @@ def _attn_fwd_inner(
|
||||
num_stages=1,
|
||||
num_warps=4,
|
||||
),
|
||||
# BI-V100 num_stages=2 variants:
|
||||
# CCCL transform benchmark (babelstream.cu) found bif=8 (64KB prefetch)
|
||||
# dominates bif=0 (32KB) on BI-V100. Physical explanation:
|
||||
# BW_per_SM × memory_latency = 56 GB/s × 1100ns ≈ 62KB
|
||||
# Triton's num_stages controls software pipelining depth, which is
|
||||
# the same concept as CCCL's bytes_in_flight. num_stages=2 doubles
|
||||
# the prefetch window, matching the 64KB sweet spot.
|
||||
# Source: cccl_upstream/cub/benchmarks/bench/transform/babelstream.cu
|
||||
triton.Config(
|
||||
{
|
||||
"BLOCK_M": 64,
|
||||
"BLOCK_N": 64,
|
||||
"waves_per_eu": 2,
|
||||
"PRE_LOAD_V": False,
|
||||
},
|
||||
num_stages=2,
|
||||
num_warps=4,
|
||||
),
|
||||
triton.Config(
|
||||
{
|
||||
"BLOCK_M": 32,
|
||||
"BLOCK_N": 64,
|
||||
"waves_per_eu": 2,
|
||||
"PRE_LOAD_V": True,
|
||||
},
|
||||
num_stages=2,
|
||||
num_warps=4,
|
||||
),
|
||||
# BI-V100 minimal tile: highest occupancy for short sequences.
|
||||
# CCCL scan benchmark found no_delay optimal (dcid=0) because
|
||||
# 16 SMs → ~32 CTAs → tile_status fits in 6MB L2 → no contention.
|
||||
# Same logic: small tiles + many CTAs maximize SM utilization.
|
||||
# Source: cccl_upstream/cub/benchmarks/bench/scan/exclusive/sum.cu
|
||||
triton.Config(
|
||||
{
|
||||
"BLOCK_M": 32,
|
||||
"BLOCK_N": 32,
|
||||
"waves_per_eu": 4,
|
||||
"PRE_LOAD_V": False,
|
||||
},
|
||||
num_stages=2,
|
||||
num_warps=2,
|
||||
),
|
||||
],
|
||||
key=['IS_CAUSAL', 'dropout_p', 'BLOCK_DMODEL'],
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user