[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:
project_6
2026-08-05 03:09:45 +00:00
parent 8a87e378f8
commit 2c43eb524f

View File

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